Switch layer still on wrong device

neural_network_pdes/euler_networks.py

@@ -60,6 +60,8 @@ def forward(self, x: Tensor) -> Tensor:
 class Net(LightningModule):
     def __init__(self, cfg: DictConfig):
         super().__init__()
+        device = "gpu" if cfg.accelerator is "cuda" else "cpu"
+
         self.automatic_optimization = False
 
         self.save_hyperparameters(cfg)
@@ -95,6 +97,7 @@ def __init__(self, cfg: DictConfig):
                 periodicity=cfg.mlp.periodicity,
                 rotations=cfg.mlp.rotations,
                 resnet=cfg.mlp.resnet,
+                device=device,
             )
         elif cfg.mlp.style == "high-order-input":
             layer_list = []
@@ -104,6 +107,7 @@ def __init__(self, cfg: DictConfig):
                 in_features=cfg.mlp.input.width,
                 out_features=cfg.mlp.hidden.width,
                 segments=cfg.mlp.input.segments,
+                device=device,
             )
             layer_list.append(input_layer)
 

pyproject.toml

@@ -13,7 +13,7 @@ hydra-nevergrad-sweeper = "1.3.0.dev0"
 omegaconf = "^2.2.3"
 Pillow = "^9.3.0"
 torch-optimizer = "^0.3.0"
-#high-order-layers-optimizers = { git = "https://github.com/jloveric/adahessian.git", branch = "master" }
+high-order-layers-optimizers = { git = "https://github.com/jloveric/adahessian.git", branch = "master" }
 grpcio = "1.51.1"
 high-order-layers-torch = "^2.0.0"