Fix prototypical_loss bug orobix#23

rcy17 · Feb 3, 2021 · 58b72c3 · 58b72c3
1 parent 5e18a5e
commit 58b72c3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/prototypical_loss.py b/src/prototypical_loss.py
@@ -81,6 +81,6 @@ def supp_idxs(c):
 
     loss_val = -log_p_y.gather(2, target_inds).squeeze().view(-1).mean()
     _, y_hat = log_p_y.max(2)
-    acc_val = y_hat.eq(target_inds.squeeze()).float().mean()
+    acc_val = y_hat.eq(target_inds.squeeze(2)).float().mean()
 
     return loss_val,  acc_val