[Bug Report] ttnn.from_torch unable to create tensor with shape (2048, 1024, 32, 32) in TILE_LAYOUT

[jerrysky3]

Describe the bug

When calling ttnn.from_torch with a torch BF16 tensor with shape (2048, 1024, 32, 32), it failed with TILE_LAYOUT (while succeeded with ROW_MAJOR_LAYOUT)

To Reproduce

Run the example below:

import torch
import ttnn

if __name__ == "__main__":
    device = ttnn.open_device(device_id=0)
    try:
        torch_x = torch.rand((2048, 1024, 32, 32), dtype=torch.bfloat16)
        x_tensor = ttnn.from_torch(torch_x, layout=ttnn.TILE_LAYOUT)
    finally:
        ttnn.close_device(device)

It failed with the error:

2024-11-14 03:39:26.736 | INFO     | SiliconDriver   - Software version 6.0.0, Ethernet FW version 6.10.0 (Device 7)
                  Metal | INFO     | Initializing device 0. Program cache is NOT enabled
                  Metal | INFO     | AI CLK for device 0 is:   1000 MHz
                  Metal | INFO     | MMIO Device 0 : Tunnel 0 : Device 0
                  Metal | INFO     | MMIO Device 0 : Tunnel 0 : Device 4
                  Metal | INFO     | Closing device 0
                  Metal | INFO     | Disabling and clearing program cache on device 0
Traceback (most recent call last):
  File "tests/ttnn/unit_tests/operations/exp_conv.py", line 128, in <module>
    x_tensor = ttnn.from_torch(torch_x, layout=ttnn.TILE_LAYOUT)
  File "/home/jerry/tt-metal/ttnn/ttnn/decorators.py", line 329, in __call__
    return self.function(*function_args, **function_kwargs)
  File "/home/jerry/tt-metal/ttnn/ttnn/operations/core.py", line 225, in from_torch
    tensor = ttnn.to_layout(tensor, layout, device=device)
  File "/home/jerry/tt-metal/ttnn/ttnn/decorators.py", line 329, in __call__
    return self.function(*function_args, **function_kwargs)
ValueError: vector
                 Device | INFO     | Closing user mode device driver

Expected behavior
The call to ttnn.from_torch should succeed

Please complete the following environment information:

• OS: Ubuntu 20.04
• Tenstorrent card: N300