add a comment

pytorch · Apr 8, 2024 · b1cc792 · b1cc792
1 parent d6b2ee6
commit b1cc792
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/ignite/metrics/mutual_information.py b/ignite/metrics/mutual_information.py
@@ -77,6 +77,7 @@ def reset(self) -> None:
 
     def _update(self, prob: torch.Tensor, log_prob: torch.Tensor) -> None:
         super()._update(prob, log_prob)
+        # We can't use += below as _sum_of_probabilities can be a scalar and prob.sum(dim=0) is a vector
         self._sum_of_probabilities = self._sum_of_probabilities + prob.sum(dim=0).to(self._device)
 
     @sync_all_reduce("_sum_of_probabilities")