def _test_ddp_cat(rank, worldsize):
setup_ddp(rank, worldsize)
dummy = Dummy()
dummy._reductions = {"foo": torch.cat}
dummy.foo = [torch.tensor([1])]
dummy._sync_dist()
assert torch.all(torch.eq(dummy.foo, torch.tensor([1, 1])))
def _test_ddp_sum(rank, worldsize):
setup_ddp(rank, worldsize)
dummy = Dummy()
dummy._reductions = {"foo": torch.sum}
dummy.foo = torch.tensor(1)
dummy._sync_dist()
assert dummy.foo == worldsize
def _test_non_contiguous_tensors(rank, worldsize):
setup_ddp(rank, worldsize)
class DummyMetric(Metric):
def __init__(self):
super().__init__()
self.add_state("x", default=[], dist_reduce_fx=None)
def update(self, x):
self.x.append(x)
def compute(self):
x = torch.cat(self.x, dim=0)
return x.sum()
metric = DummyMetric()
metric.update(torch.randn(10, 5)[:, 0])