Update test_ener_spin_model.py

iProzd · Feb 25, 2024 · 4f67f78 · 4f67f78
1 parent 6179824
commit 4f67f78
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/source/tests/pt/model/test_ener_spin_model.py b/source/tests/pt/model/test_ener_spin_model.py
@@ -49,6 +49,7 @@ class SpinTest:
     def setUp(self):
         self.prec = 1e-10
         natoms = 5
+        self.ntypes = 3  # ["O", "H", "B"] for test
         self.cell = 4.0 * torch.eye(3, dtype=dtype, device=env.DEVICE).unsqueeze(0)
         self.coord = 3.0 * torch.rand(
             [natoms, 3], dtype=dtype, device=env.DEVICE