enh: Check for Sophia before performing the Hessian update step

This makes swapping optimizers easier, see !4
token-tricksters · Jun 21, 2023 · e55e795 · e55e795
1 parent 8ce9763
commit e55e795
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/classifier.py b/classifier.py
@@ -301,7 +301,11 @@ def train(args):
 
             optimizer.zero_grad(set_to_none=True)
 
-            if iter_num % hess_interval == hess_interval - 1:
+            # Check if we use the Sophia Optimizer
+            if (
+                hasattr(optimizer, "update_hessian")
+                and iter_num % hess_interval == hess_interval - 1
+            ):
                 # Update the Hessian EMA
                 logits = model(b_ids, b_mask)
                 samp_dist = torch.distributions.Categorical(logits=logits)