ttnn.to_torch

ttnn.to_torch(tensor: ttnn.Tensor, dtype: torch.dtype = None, *, torch_rank: int = None, mesh_composer: ttnn.CppMeshToTensor = None, device: ttnn.MeshDevice = None, cq_id: int = 0) → torch.Tensor

Converts the ttnn.Tensor tensor into a torch.Tensor. It does not call to_layout for bfloat8_b or bfloat4_b as we now convert to tile layout during tensor.to_torch().

Parameters:

  • tensor (ttnn.Tensor) – the input tensor.

  • dtype (torch.dtype, optional) – the desired torch data type of returned tensor. Defaults to None.

Keyword Arguments:

  • torch_rank (int, optional) – Desired rank of the torch.Tensor. Defaults to None. Will use torch.squeeze operation to remove dimensions until the desired rank is reached. If not possible, the operation will raise an error.

  • mesh_composer (ttnn.CppMeshToTensor, optional) – The desired ttnn mesh composer. Defaults to None.

  • device (ttnn.MeshDevice, optional) – The ttnn device of the input tensor. Defaults to None.

  • cq_id (int, optional) – The command queue ID to use. Defaults to 0.

Returns:

torch.Tensor – The converted torch tensor.

Example

# Create a TT-NN tensor and convert it to a Torch tensor
ttnn_tensor = ttnn.rand((2, 3), dtype=ttnn.bfloat16, device=device)
torch_tensor = ttnn.to_torch(ttnn_tensor)

logger.info("Torch tensor shape", torch_tensor.shape)