@ptrblck, thank you for your answer!
I’ll try to exaplain why I came with that idea.
So, I’ve got a model in which I would like to have about 100 different layers as outputs (each of them has few features) (“heads”). Each of the output layer (“head”) bases on the input which is coming from the same layer “X”.
Now the tricky part: because of the problem with data, which it hasn’t got always “real” outputs to compare in loss function (with output values from the model) - I’m trying to do something like that I “turn on” only this output layers in model, to which I have data for each sample.
But now (while I was writing this) I found another problem - while I’m calculating loss from this “turned on” output layers from model - in the step loss.backward() and optimizer.step() it backpropagates through all of the output layers, not only this “turned on” for which data has passed, isn’t it? Is there possibility to backpropagate and optimize only through this output layers of model (“heads”) which are turned on while training for each sample?
Because of this “turning on” different output layers - I had to remove normalization and train with samples (bs=1). But as I think now - if I train my model using samples and without normalization, I can inference with batches (without batch normalization), can’t I?
Best regards