From 05463316372a32744534ac2e61d71d56651079e2 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Sat, 29 Aug 2026 11:18:39 +0800
Subject: [PATCH] fix: skip gradient checkpointing log when no modules
 configured

- GradientCheckpointingCallback.on_train_begin returns early on empty module list
- previously logged "Gradient checkpointing enabled" even when checkpointing was inactive, misleading profiling
---
 astrai/trainer/train_callback.py | 2 ++
 1 file changed, 2 insertions(+)

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")