Replace lightning.pytorch as pl with lightning as L

src/callbacks/metric.py

@@ -2,20 +2,20 @@
 import logging
 from pathlib import Path
 
-import lightning.pytorch as pl
+import lightning as L
 from lightning.fabric.utilities.apply_func import convert_tensors_to_scalars
 from lightning.pytorch.trainer.states import TrainerFn
 
 
-class Metric(pl.Callback):
+class Metric(L.Callback):
     r"""
     Save logged metrics to ``Trainer.log_dir``.
     """
 
     def teardown(
         self,
-        trainer: pl.Trainer,
-        pl_module: pl.LightningModule,
+        trainer: L.Trainer,
+        pl_module: L.LightningModule,
         stage: str | None = None,
     ) -> None:
         metrics = {}

src/datamodules/glue_datamodule.py

@@ -2,7 +2,7 @@
 from functools import partial
 from typing import Literal
 
-import lightning.pytorch as pl
+import lightning as L
 from datasets import load_dataset
 from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
 from torch.utils.data import DataLoader
@@ -25,7 +25,7 @@
 ]
 
 
-class GLUEDataModule(pl.LightningDataModule):
+class GLUEDataModule(L.LightningDataModule):
     task_text_field_map = {
         "cola": ["sentence"],
         "sst2": ["sentence"],

src/datamodules/mnist_datamodule.py

@@ -1,11 +1,11 @@
-import lightning.pytorch as pl
+import lightning as L
 from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
 from torch.utils.data import DataLoader, random_split
 from torchvision.datasets import MNIST
 from torchvision.transforms import transforms
 
 
-class MNISTDataModule(pl.LightningDataModule):
+class MNISTDataModule(L.LightningDataModule):
     def __init__(
         self,
         data_dir: str = "data/",

src/models/glue_transformer.py

@@ -2,7 +2,7 @@
 from typing import Any
 
 import evaluate
-import lightning.pytorch as pl
+import lightning as L
 import torch
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from transformers import (
@@ -13,7 +13,7 @@
 )
 
 
-class GLUETransformer(pl.LightningModule):
+class GLUETransformer(L.LightningModule):
     def __init__(
         self,
         task_name: str,

src/models/mnist_model.py

@@ -1,12 +1,12 @@
-import lightning.pytorch as pl
+import lightning as L
 import torch
 import torch.nn.functional as F
 from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from torchmetrics import Accuracy, MetricCollection
 
 
-class MNISTModel(pl.LightningModule):
+class MNISTModel(L.LightningModule):
     def __init__(
         self,
         input_size: int = 28 * 28,

src/utils/loggers.py

@@ -1,6 +1,6 @@
 import os
 
-import lightning.pytorch as pl
+import lightning as L
 from lightning.fabric.utilities.types import _PATH
 from lightning.pytorch.callbacks import ModelCheckpoint
 
@@ -19,10 +19,10 @@ def log_dir(self) -> str:
     return dirpath
 
 
-pl.Trainer.log_dir = log_dir
+L.Trainer.log_dir = log_dir
 
 
-def __resolve_ckpt_dir(self, trainer: pl.Trainer) -> _PATH:
+def __resolve_ckpt_dir(self, trainer: L.Trainer) -> _PATH:
     """Determines model checkpoint save directory at runtime. References attributes from the trainer's logger
     to determine where to save checkpoints. The base path for saving weights is set in this priority:
     1.  Checkpoint callback's path (if passed in)