You can directly index the tensor as seen here:
x = torch.ones(2, 2, 2)
idx = torch.tensor([0])
x[:, idx] = 0.
print(x)
> tensor([[[0., 0.],
           [1., 1.]],

          [[0., 0.],
           [1., 1.]]])