fix: correcting shape specification for target schemas

IntelLabs · Nov 27, 2024 · 7e807ca · 7e807ca
1 parent 6df90f1
commit 7e807ca
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/matsciml/datasets/tests/test_schema.py b/matsciml/datasets/tests/test_schema.py
@@ -48,10 +48,14 @@ def test_dataset_minimal_schema_pass():
         targets=[
             {
                 "name": "total_energy",
-                "shape": 0,
+                "shape": "0",
                 "description": "Total energy of the system.",
             },
-            {"name": "forces", "shape": 3, "description": "Atomic forces per node."},
+            {
+                "name": "forces",
+                "shape": "*,3",
+                "description": "Atomic forces per node.",
+            },
         ],
         split_blake2s=splits,
     )
@@ -70,10 +74,14 @@ def test_dataset_minimal_schema_roundtrip():
         targets=[
             {
                 "name": "total_energy",
-                "shape": 0,
+                "shape": "0",
                 "description": "Total energy of the system.",
             },
-            {"name": "forces", "shape": 3, "description": "Atomic forces per node."},
+            {
+                "name": "forces",
+                "shape": "*,3",
+                "description": "Atomic forces per node.",
+            },
         ],
         split_blake2s=splits,
     )