accelerator.gather() at training time

Can I use accelerator.gather() at training time? Would gradients be calculated properly? Basically my use case is something like below toy snippet. It seems that there is some issue with gradient flow in this scheme as my validation accuracy drops to 0.

model, optimizer, train_loader = accelerator.prepare(model, optimizer, train_loader)
for i, data in enumerate(train_loader):
    model.zero_grad()
    
    a, b = model(data)
    b_all = accelerator.gather(b)
    c = f(a, b_all)
    loss = criterion(a, b, c)
    accelerator.backward(loss)
    optimizer.step()

[sgugger]

Thanks for the tip! Using that it’s possible to add a gather_with_grad function that would work with training, leveraging the existing gather.

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.