Fix AGC state, param_groups

vballoli · Apr 1, 2021 · 1d8ae87 · 1d8ae87
1 parent cf39041
commit 1d8ae87
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/nfnets/agc.py b/nfnets/agc.py
@@ -50,6 +50,9 @@ def __init__(self, params, optim: optim.Optimizer, clipping: float = 1e-2, eps:
         self.agc_params = params
         self.eps = eps
         self.clipping = clipping
+
+        self.param_groups = optim.param_groups
+        self.state = optim.state
 
         #super(AGC, self).__init__([], defaults)