I think I found a way to fix the issue. We need to distinguish between an empty prediction and a non-empty prediction. Further, we need to set the gradients to none instead to of zero. Here is what I did:
x = torch.randint(0, 2, (100, 3)).float()
for t in x:
    out1, out2 = model(t.view(1, 3))

    # Store weights to compare them after updating them
    head1_weight_buffer = model.head1.weight.detach().clone()
    head2_weight_buffer = model.head2.weight.detach().clone()

    # Check where prediction is non zero and only calculate the gradients for non-zero predictions
    if out1.shape[0] > 0:
        if out2.shape[0] > 0:
            out1.mean().backward(retain_graph=True)
            out2.mean().backward()
        else:
            out1.mean().backward()
    else:
        out2.mean().backward()

    optimizer_adam.step()
    # set_to_none = True
    optimizer_adam.zero_grad(set_to_none=True)

    # Compare weights. We are expecting that just one head changed.
    print(torch.all(torch.eq(head1_weight_buffer, model.head1.weight)))
    print(torch.all(torch.eq(head2_weight_buffer, model.head2.weight)))

It solved my problem! Heads are only getting updated if we were using samples from the specific head group. Can you confirm whether the code snippet above is correct? And why do we have to distinguish between non-zero and zero prediction?