Merge pull request #110 from luigibonati/bug_unrepresented_state

Fix bug unrepresented state in TDA loss
luigibonati · Dec 22, 2023 · 9e35e07 · 9e35e07
2 parents 80824cf + dbac2c1
commit 9e35e07
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mlcolvar/core/loss/tda_loss.py b/mlcolvar/core/loss/tda_loss.py
@@ -156,7 +156,7 @@ def tda_loss(
 
     for i in range(n_states):
         # check which elements belong to class i
-        if not torch.nonzero(labels == i).any():
+        if not (labels == i).any():
             raise ValueError(
                 f"State {i} was not represented in this batch! Either use bigger batch_size or a more equilibrated dataset composition!"
             )