Update test_equivariance.py

Author: guillemsimeon

tests/test_equivariance.py

@@ -28,7 +28,7 @@ def test_scalar_invariance():
     torch.testing.assert_allclose(y, y_rot)
 
 
-@pytest.mark.parametrize("model_name", ["equivariant-transformer", "equivariant-tensornet"])
+@pytest.mark.parametrize("model_name", ["equivariant-transformer", "tensornet"])
 def test_vector_equivariance(model_name):
     torch.manual_seed(1234)
     rotate = torch.tensor(
@@ -38,14 +38,23 @@ def test_vector_equivariance(model_name):
             [-0.0626055, 0.3134752, 0.9475304],
         ]
     )
-
-    model = create_model(
-        load_example_args(
-            model_name,
-            prior_model=None,
-            output_model="VectorOutput",
+    if model_name == "equivariant_transformer"
+        model = create_model(
+            load_example_args(
+                model_name,
+                prior_model=None,
+                output_model="VectorOutput",
+            )
+        )
+    if model_name == "tensornet"
+        model = create_model(
+            load_example_args(
+                model_name,
+                prior_model=None,
+                vector_output=True,
+                output_model="VectorOutput",
+            )
         )
-    )
     z = torch.ones(100, dtype=torch.long)
     pos = torch.randn(100, 3)
     batch = torch.arange(50, dtype=torch.long).repeat_interleave(2)