ttnn.argmax fails with error "Buffer size and page size should be larger than 0 bytes"

[vigneshkeerthivasanx]

Describe the bug ttnn.argmax fails with error "Buffer size and page size should be larger than 0 bytes!" for the test case of shape [1, 50265].

To Reproduce Steps to reproduce the behavior:

1. Save the below snippet to a file

   
   import ttnn
   import torch 

def test_unit_argmax(device, reset_seeds): tensor_a = torch.rand(1, 50265) tensor_a = ttnn.from_torch(tensor_a, device = device, dtype = ttnn.bfloat16, layout = ttnn.TILE_LAYOUT) ttnn_output = ttnn.argmax(tensor_a, dim =-1)

2.Run the script using pytest `<path/to/file>`
3.See error 

Screenshots
The following error is faced when the above code is tested.

E RuntimeError: TT_FATAL @ ../tt_metal/impl/buffers/buffer.cpp:31: size != 0 and page_size != 0 E info: E Buffer size and page size should be larger than 0 bytes!

[vigneshkeerthivasanx]

cc: @saichandax @boris-drazic