add deprecation warning for old preemption callback

Signed-off-by: Ananth Subramaniam <[email protected]>
NVIDIA · Nov 21, 2024 · 8f7b01c · 8f7b01c
1 parent 8ab46ff
commit 8f7b01c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/nemo/utils/callbacks/preemption.py b/nemo/utils/callbacks/preemption.py
@@ -19,8 +19,10 @@
 from lightning.pytorch.callbacks import Callback
 
 from nemo.utils import logging
+from nemo.utils.decorators import deprecated
 
 
+@deprecated(explanation="nemo.utils.callbacks.preemption.PreemptionCallback is deprecated in favor of nemo.lightning.pytorch.callbacks.preemption.PreemptionCallback")
 class PreemptionCallback(Callback):
     """
     PreemptionCallback class creates a callback that checks for preemption during training at the end of every step.