Upgrade to PyTorch Lightning 1.8 (#1795)

* Initial commit

* Update ClampCallback for 1.8

Author: martinkim0

pyproject.toml

@@ -67,7 +67,7 @@ pyro-ppl = ">=1.6.0"
 pytest = {version = ">=4.4", optional = true}
 python = ">=3.7,<4.0"
 python-igraph = {version = "*", optional = true}
-pytorch-lightning = ">=1.7.0,<1.8"
+pytorch-lightning = ">=1.8.0,<1.9"
 rich = ">=9.1.0"
 scanpy = {version = ">=1.6", optional = true}
 scikit-learn = ">=0.21.2"

scvi/external/cellassign/_model.py

@@ -271,8 +271,8 @@ class ClampCallback(Callback):
     def __init__(self):
         super().__init__()
 
-    def on_batch_end(self, trainer, pl_module):
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
         """Clamp parameters."""
         with torch.inference_mode():
             pl_module.module.delta_log.clamp_(np.log(pl_module.module.min_delta))
-        super().on_batch_end(trainer, pl_module)
+        super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)

scvi/train/_callbacks.py

@@ -21,7 +21,7 @@ def __init__(self):
     def on_train_epoch_start(self, trainer, pl_module):
         """Subsample labels at the beginning of each epoch."""
         trainer.train_dataloader.loaders.resample_labels()
-        super().on_epoch_start(trainer, pl_module)
+        super().on_train_epoch_start(trainer, pl_module)
 
 
 class SaveBestState(Callback):
@@ -86,7 +86,7 @@ def __init__(
     def check_monitor_top(self, current):  # noqa: D102
         return self.monitor_op(current, self.best_module_metric_val)
 
-    def on_epoch_end(self, trainer, pl_module):  # noqa: D102
+    def on_val_epoch_end(self, trainer, pl_module):  # noqa: D102
         logs = trainer.callback_metrics
         self.epochs_since_last_check += 1