Ensure to index a torch.Tensor with a torch.Tensor

[FredDeCeuster]

Given the still unresolved bug in PyTorch:
pytorch/pytorch#22013
Please ensure that each torch.Tensor is indexed by a torch.Tensor.
Just to be safe, it is probably best to cast all indices always to torch.Tensor.