torch.sqeeze()

import torch

x = torch.rand(1, 1, 2, 3)

x = x.squeeze() # [1, 1, 2, 3] -> [2, 3]

print(x)

 

tensor([[0.7293, 0.5819, 0.9878],
        [0.6379, 0.6204, 0.8290]])

 

y = torch.rand(1, 1, 4)

z = y.squeeze()[2] # [1, 1, 4] -> [4]

print(z)

tensor(0.9240)