more general output for configure_optimizers

shimmer/modules/global_workspace.py

@@ -5,8 +5,11 @@
 
 import torch
 from lightning.pytorch import LightningModule
-from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfig
+from lightning.pytorch.utilities.types import (
+    OptimizerLRScheduler,
+)
 from torch.nn import Module, ModuleDict
+from torch.optim.adamw import AdamW
 from torch.optim.lr_scheduler import LRScheduler, OneCycleLR
 
 from shimmer.modules.contrastive_loss import ContrastiveLoss, ContrastiveLossType
@@ -553,15 +556,15 @@ def predict_step(  # type: ignore
         domain_latents = self.encode_domains(batch)
         return self.forward(domain_latents)
 
-    def configure_optimizers(self) -> OptimizerLRSchedulerConfig:
+    def configure_optimizers(self) -> OptimizerLRScheduler:
         """
         Configure models optimizers.
 
         Here we use `AdamW` for the optimizer and `OneCycleLR` for the learning-rate
         scheduler.
         """
 
-        optimizer = torch.optim.AdamW(
+        optimizer = AdamW(
             self.parameters(),
             lr=self.optim_lr,
             weight_decay=self.optim_weight_decay,