I have two tensors I want to update “at the same time”, and those updates depend on each other, inside a for loop. In “general Python”, I can update them on one line:
for i in range(10):
x, y = x + y/2, y + x/2
x = some_nn(x)
y = some_nn(y)
Is this safe to do in PyTorch in the forward method? Specifically is it ok for autograd, and for eventual multiple GPU training?
Alternatively, I think, I could do this as
for i in range(10):
x_old = x.clone() # possibly with .detach(), too?
y_old = y.clone()
x = x + y_old/2
y = y + x_old/2
x = some_nn(x)
y = some_nn(y)
(Of course, I could put some_function a line earlier in both cases, but I don’t think that’s important here.)
Is there a reason to choose one method over the other, or is there a better method I haven’t thought of?