added the required changes

Author: 01AbhiSingh

tests/tests_pytorch/trainer/optimization/test_optimizers.py

@@ -18,7 +18,7 @@
 
 import pytest
 import torch
-from lightning.pytorch import Trainer
+from lightning.pytorch import Trainer, LightningModule
 from lightning.pytorch.callbacks import ModelCheckpoint
 from lightning.pytorch.core.optimizer import (
     _configure_optimizers,