Merge pull request #25 from rcy17/patch-1

Fix prototypical_loss bug #23
orobix · Apr 5, 2022 · df89808 · df89808
2 parents 5e18a5e + 58b72c3
commit df89808
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