如何在不使用就地操作的情况下交换 PyTorch 张量中的值（保留梯度）

Question

I have a tensor called state of shape torch.Size([N, 2**n, 2**n]) , and I want to apply the following operations:

state[[0,1]] = state[[1,0]]
state[0] = -1*state[0]

Both of these are in-place operations. Are there some out-of-place operations that I can substitute them with? These lines are inside a for-loop , so it would be a bit difficult to just create new variables.

Answer 1

I managed to figure it out!

Replace:

state[[0,1]] = state[[1,0]] # in-place operation

with:

state = state[[1,0]] # out-of-place operation

And for the second line, we replace:

state[0] = -1*state[0] # in-place operation

with:

# out-of-place operations
temp = torch.ones(state.shape).type(state.type()).to(state.device)
temp[1] = -1*temp[1]
state = state*temp

This seems to be doing the job!