Remove adahessian since sophia is so much better

examples/function_example.py

@@ -8,7 +8,6 @@
 import numpy as np
 import torch
 import torch.optim
-import torch_optimizer as alt_optim
 from pytorch_lightning import LightningModule, Trainer
 from torch.nn import functional as F
 from torch.utils.data import DataLoader, Dataset
@@ -131,16 +130,7 @@ def train_dataloader(self):
         return DataLoader(FunctionDataset(), batch_size=4)
 
     def configure_optimizers(self):
-        if self.optimizer == "adahessian":
-            return alt_optim.Adahessian(
-                self.layer.parameters(),
-                lr=1.0,
-                betas=(0.9, 0.999),
-                eps=1e-4,
-                weight_decay=0.0,
-                hessian_power=1.0,
-            )
-        elif self.optimizer == "adam":
+        if self.optimizer == "adam":
             return torch.optim.Adam(self.parameters(), lr=0.001)
         elif self.optimizer == "lion":
             return Lion(self.parameters(), lr=0.001)

examples/invariant_mnist.py

@@ -5,7 +5,6 @@
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.optim as optim
-import torch_optimizer as alt_optim
 import torchvision
 import torchvision.transforms as transforms
 from omegaconf import DictConfig, OmegaConf
@@ -139,16 +138,7 @@ def test_step(self, batch, batch_idx):
         return self.eval_step(batch, batch_idx, "test")
 
     def configure_optimizers(self):
-        if self.cfg.optimizer.name == "adahessian":
-            return alt_optim.Adahessian(
-                self.parameters(),
-                lr=1.0,
-                betas=(0.9, 0.999),
-                eps=1e-4,
-                weight_decay=0.0,
-                hessian_power=1.0,
-            )
-        elif self.cfg.optimizer.name == "adam":
+        if self.cfg.optimizer.name == "adam":
             optimizer = optim.Adam(self.parameters(), lr=self.cfg.optimizer.lr)
             lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(
                 optimizer,

examples/variational_autoencoder.py

@@ -3,7 +3,6 @@
 import hydra
 import torch
 import torch.optim as optim
-import torch_optimizer as alt_optim
 import torchvision
 import torchvision.transforms as transforms
 from matplotlib.pyplot import figure
@@ -158,16 +157,7 @@ def test_step(self, batch, batch_idx):
         return self.eval_step(batch, batch_idx, "test")
 
     def configure_optimizers(self):
-        if self._cfg.optimizer.name == "adahessian":
-            return alt_optim.Adahessian(
-                self.parameters(),
-                lr=1.0,
-                betas=(0.9, 0.999),
-                eps=1e-4,
-                weight_decay=0.0,
-                hessian_power=1.0,
-            )
-        elif self._cfg.optimizer.name == "adam":
+        if self._cfg.optimizer.name == "adam":
 
             optimizer = optim.Adam(
                 params=self.parameters(),

pyproject.toml

@@ -1,14 +1,13 @@
 [tool.poetry]
 name = "high-order-layers-torch"
-version = "2.5.3"
+version = "2.6.0"
 description = "High order layers in pytorch"
 authors = ["jloverich <john.loverich@gmail.com>"]
 license = "MIT"
 readme = "README.md"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4"
-torch-optimizer = "^0.3.0"
 hydra-core = "^1.2.0"
 torch = ">=2.1.1"
 torchvision = ">=0.18.0"