Function MaybeSyncBatchnorm

byol_pytorch/byol_pytorch.py:41–43

(is_distributed = None)

        p.requires_grad = val

def MaybeSyncBatchnorm(is_distributed = None):
    is_distributed = default(is_distributed, dist.is_initialized() and dist.get_world_size() > 1)
    return nn.SyncBatchNorm if is_distributed else nn.BatchNorm1d

# loss fn