From 5e8a67a535abea2ff288b1b804a9aa95eba40732 Mon Sep 17 00:00:00 2001
From: David Hou <david.hou314@gmail.com>
Date: Thu, 29 Feb 2024 13:52:54 -0800
Subject: [PATCH] oops

---
 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 6fc70a1989..5bd28db7e6 100644
--- a/examples/hlb_cifar10.py
+++ b/examples/hlb_cifar10.py
@@ -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: