Thank you for your answer.
Ok, I’ll check it, but if you say so, it is probably like that and it was only my misunderstanding.

I don’t completely understand, how your decide which “heads” are used for the current input.
If you use some specific statistics for the current activation, you could probably apply these check for the complete batch and send each sample to the corresponding head.

I’m providing keys to forward pass of the model and inside I’ve got some conditions. After the training and before inferencing I will change that - I will remove “keys” in forward pass and make some condition where “turning on” output layers will be based on the output of one specific layer ( let’s say “Y”).
Example: If from layer “Y” I will get value 1, it will turn on output layer 1, 3 and 5. - and so on.
Because of the problem with data I don’t know how or even if I can implement training with batches - each sample can “turn on” different layers, so I provide to model different number of “keys” and I get different number of outputs and I can’t batch samples because input tensors (and on the other hand output tensors) have to be of the same size, haven’t they? And in my case, each sample can have different sizes of “keys” and outputs so tensors could not “merge” in batch.
I think that I can handle training with samples (bs=1) (with having increased time of training), but will it have any impact later on the inferencing with batches or it should be fine?