ttnn.squeeze

ttnn.squeeze() → None

Returns a tensor with the specified dimensions squeezed. If dim is not provided, all dimensions of size 1 will be squeezed. If dim is an integer, only the specified dimension will be squeezed. If dim is a list of integers, all specified dimensions will be squeezed.

If a specified dimension in dim does not have size 1, it will be ignored.

Equivalent pytorch code:

input_tensor = torch.rand((1,1,1,256), dtype=torch.bfloat16)
output_tensor = torch.squeeze(input_tensor, 2) # tensor of shape (1,1,256), where at dimension 2 we removed it

:param * input_tensor: Input Tensor. :param * dim: Dim where we want to squeeze