Use PL is_overridden in BYOT example (Lightning-AI#17070)

byronyi · May 5, 2023 · ea44aac · ea44aac
1 parent 5d102bf
commit ea44aac
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/examples/fabric/build_your_own_trainer/trainer.py b/examples/fabric/build_your_own_trainer/trainer.py
@@ -4,14 +4,15 @@
 from typing import Any, cast, Iterable, List, Literal, Optional, Tuple, Union
 
 import torch
-from lightning_utilities import apply_to_collection, is_overridden
+from lightning_utilities import apply_to_collection
 from tqdm import tqdm
 
 import lightning as L
 from lightning.fabric.accelerators import Accelerator
 from lightning.fabric.loggers import Logger
 from lightning.fabric.strategies import Strategy
 from lightning.fabric.wrappers import _unwrap_objects
+from lightning.pytorch.utilities.model_helpers import is_overridden
 
 
 class MyCustomTrainer:
@@ -274,7 +275,7 @@ def val_loop(
             return
 
         # no validation but warning if val_loader was passed, but validation_step not implemented
-        if val_loader is not None and not is_overridden("validation_step", _unwrap_objects(model), L.LightningModule):
+        if val_loader is not None and not is_overridden("validation_step", _unwrap_objects(model)):
             L.fabric.utilities.rank_zero_warn(
                 "Your LightningModule does not have a validation_step implemented, "
                 "but you passed a validation dataloder. Skipping Validation."