#0: correct fallback to not throw when tensor is not on device

tenstorrent · May 17, 2024 · c1afeeb · c1afeeb
1 parent 9831fa9
commit c1afeeb
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/ttnn/ttnn/operations/fallback.py b/ttnn/ttnn/operations/fallback.py
@@ -5,12 +5,14 @@
 import sys
 import ttnn
 import torch
+import tt_lib
 
 THIS_MODULE = sys.modules[__name__]
 
 
 def reshape(tensor: ttnn.Tensor, shape) -> ttnn.Tensor:
+    device = tensor.device() if tensor.storage_type() == tt_lib.tensor.StorageType.DEVICE else None
     torch_tensor = ttnn.to_torch(tensor)
     torch_tensor = torch.reshape(torch_tensor, tuple(shape))
-    tensor = ttnn.from_torch(torch_tensor, layout=tensor.layout, dtype=tensor.dtype, device=tensor.device())
+    tensor = ttnn.from_torch(torch_tensor, layout=tensor.layout, dtype=tensor.dtype, device=device)
     return tensor