Hi,
I created one model composed of two “parts” (though it is wrapped as one model). I would like to decompose this model into two models because I want to train each model with different batch sizes. So for example, the first model/part would use a batch size of 64, and the second/part one would use a batch size of 128.
For the forward pass, there is no issue because I can run the “forward” method for the first model by decomposing it into batches. The issue is for the gradient. The autograd in Pytorch seems to compute (and allocate memory) according to the batch size of the second model, and I do not know how to “tell” Pytorch to treat the first part with a different batch size. Hence, it is allocating way too much memory than I would like it to.
Thanks for the help!