From e0fd41bb463ef9d9df3e281e770ee4fc23bb56e5 Mon Sep 17 00:00:00 2001 From: David Hou Date: Thu, 29 Feb 2024 12:12:47 -0800 Subject: [PATCH] manually shard running mean and running var --- examples/hlb_cifar10.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/hlb_cifar10.py b/examples/hlb_cifar10.py index 0e1b29d9a2..cc62b4bac3 100644 --- a/examples/hlb_cifar10.py +++ b/examples/hlb_cifar10.py @@ -305,8 +305,11 @@ def train_cifar(): X_test, Y_test = X_test.cast(dtypes.default_float), Y_test.cast(dtypes.default_float) if len(GPUS) > 1: - for x in get_parameters(model): - x.to_(GPUS) + for k, x in get_state_dict(model): + if 'running_mean' in k or 'running_bias' in k: + x.shard_(GPUS, axis=0) + else: + x.to_(GPUS) # parse the training params into bias and non-bias params_dict = get_state_dict(model)