sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

oops

Browse Source
This commit is contained in:
David Hou
2024-02-29 13:52:54 -08:00
parent 7ebf65d89a
commit 5e8a67a535
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/hlb_cifar10.py
+1 -1
View File
@@ -262,7 +262,7 @@ def train_cifar():
X_test, Y_test = X_test.cast(dtypes.default_float), Y_test.cast(dtypes.default_float)
if len(GPUS) > 1:
for k, x in get_state_dict(model):
for k, x in get_state_dict(model).items():
if not getenv('SYNCBN') and ('running_mean' in k or 'running_bias' in k):
x.shard_(GPUS, axis=0)
else: