Accessing arbitrary dimension

I have a N-D tensor where N is arbitrary. What I want to do is I want to access the last dimension of the tensor. For example, say I have a tensor of shape [A, B, C, K], I want to is something like

a = tc.randn(A, B, C, K)
a[:, :, :, K] = 1

except that there can be any number of dimensions before K. It seems that if I only want to get the value I can use index_select, but index_select doesn’t share the same storage with the original tensor so I can’t modify the values using it.

Would tensor.index_fill_ work?

1 Like

Thanks, that solved my problem. I thought that tensor operations are all available under torch so I was not looking for the method in torch.Tensor :joy:

Another follow-up question: how should I set it to another tensor? For example, instead of

a[:, :, :, K]=1

what should I do to achieve

o = tc.randn(A, B, C)
a[:, :, :, K] = o

It seems that torch.narrow works in this case.