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.