Fix AGC model bug

vballoli · Mar 9, 2021 · ad2dbdc · ad2dbdc
1 parent 5dcdba6
commit ad2dbdc
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/nfnets/agc.py b/nfnets/agc.py
@@ -45,8 +45,7 @@ def __init__(self, params, optim: optim.Optimizer, clipping: float = 1e-2, eps:
                           module in model.named_modules() if name not in ignore_agc]
 
         else:
-            params = [{"params": list(module.parameters())} for name,
-                          module in model.named_modules()]
+            params = [{"params": params}]
 
         self.agc_params = params
         self.eps = eps

diff --git a/setup.py b/setup.py
@@ -8,7 +8,7 @@
 setup(
   name = 'nfnets-pytorch',
   packages = find_packages(),
-  version = '0.1.1',
+  version = '0.1.2',
   license='MIT',
   description = 'NFNets, PyTorch',
   long_description=long_description,

diff --git a/tests/test_agc.py b/tests/test_agc.py
@@ -4,11 +4,20 @@
 
 from nfnets import replace_conv, AGC
 
-def test_agc():
+def test_agc_model():
   model = resnet18()
   replace_conv(model)
   optim = SGD(model.parameters(), 1e-3)
   optim = AGC(model.parameters(), optim, model=model)
   optim.zero_grad()
   model(torch.randn(1,3,64,64)).sum().backward()
+  optim.step()
+
+def test_agc():
+  model = resnet18()
+  replace_conv(model)
+  optim = SGD(model.parameters(), 1e-3)
+  optim = AGC(model.parameters(), optim)
+  optim.zero_grad()
+  model(torch.randn(1,3,64,64)).sum().backward()
   optim.step()