Update app.py

app.py

@@ -60,11 +60,12 @@ if text:
             return_tensors='pt',
             return_length=True
         )
-        input_ids = encoded_dict['length'].unsqueeze(0)
-        attn_mask = torch.arange(input_ids.size(1)).to(device)
+        input_ids = encoded_dict['input_ids']
+        input_ids_len = encoded_dict['length'].unsqueeze(0)
+        attn_mask = torch.arange(input_ids.size(1))
         attn_mask = attn_mask[None, :] < input_ids_len[:, None]
-        outputs = model(encoded_dict['input_ids'], attn_mask)
+        outputs = model(input_ids=input_ids, attention_mask=attn_mask)
 
-        _, preds = torch.max(outputs, 1)
+        _, preds = torch.max(outputs.logits, 1)
 
     st.write(topics_raw[preds.squeeze(0)])