def test_sync_dist(_): value = LightningModule._LightningModule__sync( torch.tensor([1.0]), sync_fn=TPUSpawnPlugin().reduce, sync_dist=True, sync_dist_op=torch.distributed.ReduceOp.SUM) assert value.item() == 8
def _ddp_test_fn(rank, worldsize): _setup_ddp(rank, worldsize) tensor = torch.tensor([1.0]) actual = LightningModule._LightningModule__sync( tensor, sync_dist=True, sync_dist_op=torch.distributed.ReduceOp.SUM) assert actual.item() == dist.get_world_size( ), "Result-Log does not work properly with DDP and Tensors"