From 1d8ae8708eee1954b08b29ce9324bdb2f91151ba Mon Sep 17 00:00:00 2001
From: vballoli <balloli.vb@gmail.com>
Date: Thu, 1 Apr 2021 02:32:45 +0200
Subject: [PATCH] Fix AGC state, param_groups

---
 nfnets/agc.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/nfnets/agc.py b/nfnets/agc.py
index a1cd95a..5867220 100644
--- 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)