Hi,
I think your main worry is what will happen if you forward multiple times on the same nn.Module.
Pytorch will create and keep around all the necessary Tensors to be able to compute the gradients corresponding to what you did in the forward. In particular, if you forward twice the same Linear layer, it will keep around both inputs and weights to be able to compute all the gradients.
This has a quite large memory impact. You can check for example the memory usage just after creating/loading your model and just after doing one forward pass. You will see that the memory consumption increase significantly (more than double in most cases) after the first forward pass. But this is expected and you can’t really get around it (see torch.utils.checkpoint if you have memory issues  )