From 99536555e45b35b87b04898205a23df35b60b239 Mon Sep 17 00:00:00 2001 From: David Hou Date: Thu, 22 Feb 2024 14:51:16 -0800 Subject: [PATCH] pass num_devices to UnsyncedBatchNorm in test, allow UnsyncedBatchNorm to be used with LB --- examples/hlb_cifar10.py | 2 +- test/test_multitensor.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/hlb_cifar10.py b/examples/hlb_cifar10.py index e880da6603..0e1b29d9a2 100644 --- a/examples/hlb_cifar10.py +++ b/examples/hlb_cifar10.py @@ -40,7 +40,7 @@ class UnsyncedBatchNorm: self.num_batches_tracked = Tensor.zeros(1, requires_grad=False) def __call__(self, x:Tensor): - assert isinstance(x.lazydata, MultiLazyBuffer) and len(x.lazydata.lbs) == self.num_devices and x.lazydata.axis == 0 + if isinstance(x.lazydata, MultiLazyBuffer): assert x.lazydata.axis is None or x.lazydata.axis == 0 and len(x.lazydata.lbs) == self.num_devices rshape, x = x.shape, x.reshape(self.num_devices, -1, *x.shape[1:]) batch_mean, batch_invstd = self.calc_stats(x) diff --git a/test/test_multitensor.py b/test/test_multitensor.py index 77b0bbd218..4cfa60bd5e 100644 --- a/test/test_multitensor.py +++ b/test/test_multitensor.py @@ -582,7 +582,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase): with Tensor.train(): synced_bn = BatchNorm2d(8) - unsynced_bn = UnsyncedBatchNorm(8) + unsynced_bn = UnsyncedBatchNorm(8, num_devices=len(devices)) for p in get_parameters([synced_bn, unsynced_bn]): p.shard_(devices)