option to use Adam optim as in deepgo

Author: aditya0by0

chebai/models/ffn.py

@@ -14,10 +14,14 @@ def __init__(
         hidden_layers: List[int] = [
             1024,
         ],
+        use_adam_optimizer: bool = False,
         **kwargs,
     ):
         super().__init__(**kwargs)
 
+        self.use_adam_optimizer: bool = bool(use_adam_optimizer)
+        print(f"Using Adam optimizer: {self.use_adam_optimizer}")
+
         layers = []
         current_layer_input_size = self.input_dim
         for hidden_dim in hidden_layers:
@@ -62,6 +66,21 @@ def forward(self, data, **kwargs):
         x = data["features"]
         return {"logits": self.model(x)}
 
+    def configure_optimizers(self, **kwargs) -> torch.optim.Optimizer:
+        """
+        Configures the optimizers.
+
+        Args:
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            torch.optim.Optimizer: The optimizer.
+        """
+        if self.use_adam_optimizer:
+            return torch.optim.Adam(self.parameters(), **self.optimizer_kwargs)
+
+        return torch.optim.Adamax(self.parameters(), **self.optimizer_kwargs)
+
 
 class Residual(nn.Module):
     """