fix ndarray behavior dict_to_array

ACEsuit · Nov 1, 2024 · 74ecfda · 74ecfda
1 parent 0764dab
commit 74ecfda
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/mace/tools/scripts_utils.py b/mace/tools/scripts_utils.py
@@ -627,7 +627,9 @@ def custom_key(key):
 
 
 def dict_to_array(input_data, heads):
-    if not all(isinstance(value, dict) for value in input_data.values()):
+    if all(isinstance(value, np.ndarray) for value in input_data.values()):
+        return np.array([input_data[head] for head in heads])
+    elif not all(isinstance(value, dict) for value in input_data.values()):
         return np.array([[input_data[head]] for head in heads])
     unique_keys = set()
     for inner_dict in input_data.values():