update: recipe

Author: kozistr

tests/constants.py

@@ -408,4 +408,5 @@
     (AdaMod, {'lr': 1e2, 'weight_decay': 1e-3, 'adam_debias': True}, 20),
     (AdaMax, {'lr': 1e0, 'weight_decay': 1e-3, 'adam_debias': True}, 5),
     (AvaGrad, {'lr': 1e1, 'weight_decay': 1e-3, 'adam_debias': True}, 5),
+    (AdaHessian, {'lr': 5e0, 'weight_decay': 1e-3, 'adam_debias': True}, 5),
 ]

tests/test_optimizers.py

@@ -66,10 +66,7 @@ def _closure() -> float:
         if init_loss == np.inf:
             init_loss = loss
 
-        if optimizer_name in ('AdaHessian', 'SophiaH'):
-            loss.backward(create_graph=True)
-        else:
-            loss.backward()
+        loss.backward(create_graph=optimizer_name in ('AdaHessian', 'SophiaH'))
 
         optimizer.step(closure(loss) if optimizer_name == 'AliG' else None)
 
@@ -236,7 +233,7 @@ def test_adamd_optimizers(optimizer_config, environment):
         if init_loss == np.inf:
             init_loss = loss
 
-        loss.backward()
+        loss.backward(create_graph=optimizer_class.__name__ in ('AdaHessian',))
 
         optimizer.step()