Yes it will create a node in your graph. However, using for loops in a model is problematic and very slow. Check the discussion here: Loops inside a model slows gradient calculation dramatically - autograd - PyTorch Forums I think you can try reshaping the vector to get the spatial form you want. Btw you are not defining the n_batches variable anywhere in the Model class so I am not sure what do you mean by that.
1 Like