fix edge case (#305)

owkin · Jun 12, 2024 · c7478b9 · c7478b9
1 parent 8799494
commit c7478b9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flamby/datasets/fed_synthetic/loss.py b/flamby/datasets/fed_synthetic/loss.py
@@ -8,4 +8,4 @@ def __init__(self):
         self.ce = torch.nn.CrossEntropyLoss()
 
     def forward(self, input: torch.Tensor, target: torch.Tensor):
-        return self.ce(input, target.squeeze().long())
+        return self.ce(input, target.squeeze(axis=1).long())