You are right, I can see that the gradients are None after the first backward call. However, in some cases where the batch does only contain samples from one group, both heads are updated. I’m expecting that in such a case only the head gets updated which the group belongs to.
I noticed that in the first rounds of training the behavior is as expected. For example, imagine the following batches:

bs1:= x=1, y=1
bs2:= x=1, y=0
bs3:= x=0, y=1

Here, in the first training loop (using only bs1), only head1 is updated. The same applies for the second training loop. However, for the third training loop, both heads are updated even tough the gradients are zero for prediction1.
I modified your example such that you can convince yourself:
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()

    for name, param in model.named_parameters():
         print(name, param.grad)

    out1.mean().backward(retain_graph=True)
    for name, param in model.named_parameters():
        print(name, param.grad)

    out2.mean().backward()
    for name, param in model.named_parameters():
        print(name, param.grad)

    optimizer_adam.step()
    optimizer_adam.zero_grad()

    # 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)))

I’m expecting that one of the print statements is true and the other one is false, but they are never false or true at the same time.
Regarding your index question: It looks like your approach as well as my approach do work. I indeed got separate sets of input data. The only difference is the idx.flatten() instead of idx[:, 0], right?