From 106ac4be14342af62b5f34df69b2fd3f85ce267e Mon Sep 17 00:00:00 2001 From: David Hou Date: Thu, 29 Feb 2024 12:20:11 -0800 Subject: [PATCH] don't shard unless syncbn=0 --- examples/hlb_cifar10.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/hlb_cifar10.py b/examples/hlb_cifar10.py index cc62b4bac3..2e3760cdd3 100644 --- a/examples/hlb_cifar10.py +++ b/examples/hlb_cifar10.py @@ -306,7 +306,7 @@ def train_cifar(): if len(GPUS) > 1: for k, x in get_state_dict(model): - if 'running_mean' in k or 'running_bias' in k: + if not getenv('SYNCBN') and ('running_mean' in k or 'running_bias' in k): x.shard_(GPUS, axis=0) else: x.to_(GPUS)