Fixed out_features with slicing option

luigibonati · May 8, 2024 · 25b0ff8 · 25b0ff8
1 parent 78ada65
commit 25b0ff8
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mlcolvar/core/transform/descriptors/pairwise_distances.py b/mlcolvar/core/transform/descriptors/pairwise_distances.py
@@ -38,7 +38,10 @@ def __init__(self,
         torch.Tensor
             Non duplicated pairwise distances between all the atoms
         """
-        super().__init__(in_features=int(n_atoms*3), out_features=int(n_atoms*(n_atoms-1) / 2))
+        if slicing_pairs is None:
+            super().__init__(in_features=int(n_atoms*3), out_features=int(n_atoms*(n_atoms-1) / 2))
+        else: 
+            super().__init__(in_features=int(n_atoms*3), out_features=len(slicing_pairs))
 
         # parse args
         self.n_atoms = n_atoms