Update accessibility_classifier/classifier/model.py

accessibility_classifier/classifier/model.py

@@ -92,10 +92,8 @@ class Model:
         with torch.no_grad():
             probabilities = self.classifier(input_ids, attention_mask)
 
-        prediction = F.softmax(probabilities.logits,
-                               dim=1).cpu().numpy().flatten().max()
-        prediction_index = np.where(F.softmax(probabilities.logits,
-                                              dim=1).cpu().numpy() == prediction)[1][0]
+        prediction = F.softmax(probabilities.logits, dim=1).cpu().numpy().flatten().max()
+        prediction_index = np.where(F.softmax(probabilities.logits, dim=1).cpu().numpy() == prediction)[1][0]
         label = self.labels[prediction_index]
 
         all_predictions = F.softmax(