I have a zero tensor in the forward pass of my module, whose size is depending on the batch size. Currently i am creating on each pass the tensor and sending it to the GPU if the input is on the GPU.
Can i somehow put this zero tensor once on the GPU and clear it after each step?
Easiest way would be to declare it as a parameter of my module and switch off requires_gradient.
I checked out : Dynamic parameter declaration in forward function 16 which seems to be what i want, but i am unable to send the tensor to the GPU. The tensor does not require a gradient, it is only there to collect the output.