add numpy target support

materialsvirtuallab · May 20, 2019 · 57aee9f · 57aee9f
1 parent f8bbc43
commit 57aee9f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/megnet/utils/general_utils.py b/megnet/utils/general_utils.py
@@ -7,6 +7,8 @@ def to_list(x):
     """
     if isinstance(x, list) or isinstance(x, tuple):
         return x
+    elif isinstance(x, np.ndarray):
+        return to_list(x.tolist())
     else:
         return [x]
 
@@ -20,4 +22,4 @@ def expand_1st(x):
     Returns:
          (np.array)
     """
-    return np.expand_dims(x, axis=0)
+    return np.expand_dims(x, axis=0)