fix shape

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
deepmodeling · Jul 22, 2024 · 29c1b58 · 29c1b58
1 parent fa94e8e
commit 29c1b58
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py
@@ -167,8 +167,11 @@ def process_spin_output_lower(
         )
         if add_mag:
             extended_out_real = extended_out_real + extended_out_mag
+        shape2 = 1
+        for ss in extended_out_tensor.shape[2:]:
+            shape2 *= ss
         extended_out_mag = (
-            extended_out_mag.view([nframes, nall, -1]) * atomic_mask
+            extended_out_mag.view([nframes, nall, shape2]) * atomic_mask
         ).view(extended_out_mag.shape)
         return extended_out_real, extended_out_mag, atomic_mask > 0.0