Excellent! This is definitely the right approach.
I suppose this is more of a python question, but with this exact approach, the training loop exits on the first StopIteration, which is a problem if the two dataset sizes differ. Is the only way to handle this in the except block?
Here’s a simulation of that situation from your code block:
dataset_head1 = TensorDataset(torch.randn(10, 1), torch.randn(10, 1))
dataset_head2 = TensorDataset(torch.randn(2, 1), torch.randn(2, 1))

for epoch in range(2):
    print('epoch ', epoch)
    # in epoch loop
    loader1 = DataLoader(dataset_head1)
    loader2 = DataLoader(dataset_head2)
    
    iter_loader1 = iter(loader1)  # length - 10
    iter_loader2 = iter(loader2)  # length - 2
    
    try:
        while True:
            # train head1
            data, target = next(iter_loader1)
            print('training head1')
    
            # train head2
            data, target = next(iter_loader2)
            print('training head2')
    except StopIteration:
        # train remaining iterations from `iter_loader1` here?
        # Use `DataLoader`._num_yielded?