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