Is there a way to alter a tensor given an input dimension and indexes? As an example, suppose I wanted to multiply the first and second index of the third dimension of a tensor by 2. I could do:

Perhaps there is a better way, but here’s an ugly solution involving transposing the desired dimension into the first position, making the change there, and then transposing it back.

t = torch.randn([4, 4, 4]) # our starting tensor
# your way
t0 = torch.clone(t)
t0[:, [1, 2]] = 1.
# proposed way, via transposing
dim = 1 # desired dimension to edit
t1 = torch.clone(t)
dims = list(range(t1.ndim))
dims.remove(dim)
dims = [dim] + dims
t1 = t1.permute(dims)
t1[[1, 2]] = 1
t1 = t1.permute(tuple(np.argsort(dims)))
print(torch.allclose(t0, t1))
Output:
True