Yes the package can do that. However, tensor needs to be of same shape and size in all processes. Then you can do something like:
for iteration, data0, data1 in enumerate(data_loader, start_iter):
    tensor = model(data0)
    synchronize()  # You probably do not need this since all_gather will force a sync
    gather_list = [torch.empty_like(tensor) for i in range(dist.get_world_size())]
    gather_list = diffdist.functional.all_gather(gather_list, tensor)
    loss = model(data1, gather_list)

Keep in mind though that all_gather is not very fast because its backprop involves running dist.reduce multiple times. When pytorch adds support for reduce_scatter, I will update the package to speed up the backprop.