Make mlc.torch.tonp() gracefully handle non-Tensors

mxbi · Dec 31, 2018 · b8a2dda · b8a2dda
1 parent b8f6f3e
commit b8a2dda
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/mlcrate/torch.py b/mlcrate/torch.py
@@ -12,8 +12,13 @@ def _check_torch_import():
 
 def tonp(tensor):
     """Takes any PyTorch tensor and converts it to a numpy array or scalar as appropiate.
+    When given something that isn't a PyTorch tensor, it will attempt to convert to a NumPy array or scalar anyway.
     Not heavily optimized."""
-    arr = tensor.data.detach().cpu().numpy()
+    _check_torch_import()
+    if isinstance(tensor, torch.Tensor):
+        arr = tensor.data.detach().cpu().numpy()
+    else: # It's not a tensor! We'll handle it anyway
+        arr = np.array(tensor)
     if arr.shape == ():
         return np.asscalar(arr)
     else: