Merge branch 'main' into ensemble_model_support

Author: RaulPPelaez

tests/test_model.py

@@ -116,28 +116,26 @@ def test_cuda_graph_compatible(model_name):
     if not torch.cuda.is_available():
         pytest.skip("CUDA not available")
     z, pos, batch = create_example_batch()
-    args = {
-        "model": model_name,
-        "embedding_dimension": 128,
-        "num_layers": 2,
-        "num_rbf": 32,
-        "rbf_type": "expnorm",
-        "trainable_rbf": False,
-        "activation": "silu",
-        "cutoff_lower": 0.0,
-        "cutoff_upper": 5.0,
-        "max_z": 100,
-        "max_num_neighbors": 128,
-        "equivariance_invariance_group": "O(3)",
-        "prior_model": None,
-        "atom_filter": -1,
-        "derivative": True,
-        "check_error": False,
-        "static_shapes": True,
-        "output_model": "Scalar",
-        "reduce_op": "sum",
-        "precision": 32,
-    }
+    args = {"model": model_name,
+            "embedding_dimension": 128,
+            "num_layers": 2,
+            "num_rbf": 32,
+            "rbf_type": "expnorm",
+            "trainable_rbf": False,
+            "activation": "silu",
+            "cutoff_lower": 0.0,
+            "cutoff_upper": 5.0,
+            "max_z": 100,
+            "max_num_neighbors": 128,
+            "equivariance_invariance_group": "O(3)",
+            "prior_model": None,
+            "atom_filter": -1,
+            "derivative": True,
+            "check_errors": False,
+            "static_shapes": True,
+            "output_model": "Scalar",
+            "reduce_op": "sum",
+            "precision": 32 }
     model = create_model(args).to(device="cuda")
     model.eval()
     z = z.to("cuda")