That sounds reasonable! However, without the if else statement, both heads are again updated, even though I’m using only an example from one of the groups. I think that is because of the retain_graph = True parameter in backward(). In the case where out1 is zero, I retain the graph which leads (somehow?) to an update of the weights in head1 and head2. Thank you very much for your help!