Improve the TensorNet model class coverage (#276)

* improve TensorNet model coverage * Update pyproject.toml Signed-off-by: Tsz Wai Ko <[email protected]> * Improve the unit test for SO(3) equivarance in TensorNet class --------- Signed-off-by: Tsz Wai Ko <[email protected]>
materialsvirtuallab · Jun 22, 2024 · cfd4874 · cfd4874
1 parent c612f62
commit cfd4874
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/models/test_tensornet.py b/tests/models/test_tensornet.py
@@ -21,6 +21,8 @@ def test_model(self, graph_MoS):
         os.remove("model.pt")
         os.remove("model.json")
         os.remove("state.pt")
+        model = TensorNet(is_intensive=False, equivariance_invariance_group="SO(3)")
+        assert torch.numel(output) == 1
 
     def test_exceptions(self):
         with pytest.raises(ValueError, match="Invalid activation type"):