Thanks for a quick reply. I really counted you will reply ptrblck 



 ptrblck:


You pass (some) parameters to the optimizer, which are updated using their gradients (stored in the .grad attribute of each parameter). If you don’t pass certain parameters to an optimizer, they won’t be updated in any way. Since softmax does not contain any parameters, your current idea won’t work unfortunately.



Maybe I did not explain this part well enough or I do not understand your reply properly.
Moreover, maybe having 3 datasets with 50,000 entities each in my example was wrong decision. Datasets may have any number of entities, so the size of the softmax layers can vary a lot - that’s why using multiple softmax layers seems a must.
Of course softmax has no parameters, so I am not willing to update softmax layer with backprop. As I assumed that optimizers are ‘binded’ to all modules in the model, so I reckoned that for each model version (that are different only in softmax layer) should have a separate optimizer, as the model would be different due to different softmax layers.
However, after reading your reply it occurred to me that maybe there is an easier approach. As softmax has no parameters, can I have optimizer binded to all layers of the model but for the softmax layer that would optimize all models’ versions parameters? So one optimizer for all models with different softmax layer?



 ptrblck:


If a layer is not used during the forward pass, it won’t get any gradients in the backward pass. You could therefore use an if condition and use different layers in the forward pass. Only the selected (and used) layer will be included in the backward pass. This will add some overhead, but you could transfer all unused layers back to the CPU to free the GPU memory for the current one.



This gives me a bit of hope, but I am confused how the GPU memory can be freed with this ‘froward approach’. If the model is loaded into GPU, all softmax layers will be sitting there as well all the time, even if they are not used in the current forward pass, am I right?
The problem is, that during the forward pass only softmax layer would be changing, which does not have any parameters anyway. I can’t see any savings here.



 ptrblck:

I would generally “Go for it!”, especially after reading your first sentence 


Thanks! I will try, maybe asking a couple of more questions