From dcfe4ac3c05922109c806bc90165121261e913a1 Mon Sep 17 00:00:00 2001 From: Stuti Raizada Date: Wed, 28 Aug 2024 09:35:07 -0500 Subject: [PATCH] #0: unitlize tensor on device --- ttnn/ttnn/operations/core.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ttnn/ttnn/operations/core.py b/ttnn/ttnn/operations/core.py index 4e1ca5006f8..575c8b17fc8 100644 --- a/ttnn/ttnn/operations/core.py +++ b/ttnn/ttnn/operations/core.py @@ -309,12 +309,12 @@ def to_torch( tensor([[-0.3008, -0.8438, 0.3242], [ 0.9023, -0.5820, 0.5312]], dtype=torch.bfloat16) """ + if tensor.layout != ttnn.ROW_MAJOR_LAYOUT: + tensor = ttnn.to_layout(tensor, ttnn.ROW_MAJOR_LAYOUT) + if ttnn.is_tensor_storage_on_device(tensor): tensor = ttnn.from_device(tensor, cq_id=cq_id) - if tensor.layout != ttnn.ROW_MAJOR_LAYOUT: - tensor = tensor.to(ttnn.ROW_MAJOR_LAYOUT, device) - shape_without_tile_padding = tuple(tensor.shape) if tensor.storage_type() == ttnn.DEVICE_STORAGE_TYPE: raise RuntimeError("ttnn.Tensor cannot be on device when converting to torch.Tensor!")