[fix] cleaning

bioimage_embed/lightning/pyro.py

@@ -7,6 +7,10 @@
 
 
 class LitAutoEncoderPyro(pl.LightningModule):
+    """
+    WIP Unsupported
+    """
+
     def __init__(self, model, batch_size=1, learning_rate=1e-3):
         super().__init__()
         # self.autoencoder = AutoEncoder(batch_size, 1)
@@ -59,6 +63,3 @@ def pyro_training_step(self, train_batch, batch_idx):
 
     def training_step(self, train_batch, batch_idx):
         return self.torch_training_step(train_batch, batch_idx)
-
-    def training_step(self, train_batch, batch_idx):
-        return self.pyro_training_step(train_batch, batch_idx)

bioimage_embed/lightning/torch.py

@@ -143,16 +143,6 @@ def eval_step(self, batch, batch_idx):
         """
         return self.predict_step(batch, batch_idx)
 
-    # def lr_scheduler_step(self, epoch, batch_idx, optimizer, optimizer_idx, second_order_closure=None):
-    #     # Implement your own logic for updating the lr scheduler
-    #     # This method will be called at each training step
-    #     # Update the lr scheduler based on the provided arguments
-    #     # You can access the lr scheduler using `self.lr_schedulers()`
-
-    #     # Example:
-    #     for lr_scheduler in self.lr_schedulers():
-    #         lr_scheduler.step()
-
     def timm_optimizers(self, model):
         optimizer = optim.create_optimizer(self.args, model.parameters())
         lr_scheduler = scheduler.create_scheduler(self.args, optimizer)[0]
@@ -168,8 +158,6 @@ def timm_to_lightning(self, optimizer, lr_scheduler):
         }
 
     def configure_optimizers(self):
-        # optimizer = optim.create_optimizer(self.args, self.model.parameters())
-        # lr_scheduler = scheduler.create_scheduler(self.args, optimizer)[0]
         optimizer, lr_scheduler = self.timm_optimizers(self.model)
         return self.timm_to_lightning(optimizer, lr_scheduler)