Update app.py

app.py

@@ -60,9 +60,10 @@ if text:
             return_tensors='pt',
             return_length=True
         )
-        
-        
-        outputs = model(encoded_dict['input_ids'], encoded_dict['length'].unsqueeze(0))
+        input_ids = encoded_dict['length'].unsqueeze(0)
+        attn_mask = torch.arange(input_ids.size(1)).to(device)
+        attn_mask = attn_mask[None, :] < input_ids_len[:, None]
+        outputs = model(encoded_dict['input_ids'], attn_mask)
 
         _, preds = torch.max(outputs, 1)