Small fix for prediction calculation for bert models

advaithsrao · Nov 26, 2023 · e8dc4a0 · e8dc4a0
1 parent 38060b7
commit e8dc4a0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/detector/modeler.py b/detector/modeler.py
@@ -264,7 +264,7 @@ def predict(
             with torch.no_grad():
                 outputs = self.model(b_input_ids, attention_mask=b_input_mask)
                 loss = outputs[0]
-                logits = F.softmax(outputs[1], dim=1)   # Taking the softmax of output
+                logits = F.softmax(outputs[1])   # Taking the softmax of output
 
             _, prediction= torch.max(logits, dim=1)
 
@@ -552,7 +552,7 @@ def predict(
             with torch.no_grad():
                 outputs = self.model(b_input_ids, attention_mask=b_input_mask)
                 loss = outputs[0]
-                logits = F.softmax(outputs[1], dim=1)   # Taking the softmax of output
+                logits = F.softmax(outputs[1])   # Taking the softmax of output
 
             _, prediction= torch.max(logits, dim=1)