allow for model to auto norm weights on optimizer step

Author: lucidrains

nGPT_pytorch/nGPT.py

@@ -459,6 +459,14 @@ def norm_weights_(self):
 
             module.norm_weights_()
 
+    def register_step_post_hook(self, optimizer):
+        assert hasattr(optimizer, 'register_step_post_hook')
+
+        def hook(*_):
+            self.norm_weights_()
+
+        return optimizer.register_step_post_hook(hook)
+
     def forward(
         self,
         ids,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nGPT-pytorch"
-version = "0.1.19"
+version = "0.1.20"
 description = "nGPT"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

train.py

@@ -139,6 +139,11 @@ def __getitem__(self, index):
 train_loader = cycle(train_loader)
 val_loader = cycle(val_loader)
 
+# if not using parametrize, register normalizing on optimizer step
+
+if not USE_PARAMETRIZE:
+    model.register_step_post_hook(optim)
+
 # training
 
 for i in tqdm.tqdm(range(NUM_BATCHES), mininterval = 10.0, desc = "training"):
@@ -159,9 +164,6 @@ def __getitem__(self, index):
 
     optim.zero_grad()
 
-    if not USE_PARAMETRIZE:
-        model.norm_weights_()
-
     if i % VALIDATE_EVERY == 0:
         model.eval()
         with torch.no_grad():