diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py
index f5fe057..3149280 100644
--- a/astrai/trainer/train_callback.py
+++ b/astrai/trainer/train_callback.py
@@ -116,6 +116,8 @@ class GradientCheckpointingCallback(TrainCallback):
             del module._original_forward
 
     def on_train_begin(self, context: TrainContext):
+        if not self.modules:
+            return
         context.model.apply(self._enable)
         logger.info("Gradient checkpointing enabled")