From 15f85a81ee7b8d4369bf0058fbc93e8bd7c69118 Mon Sep 17 00:00:00 2001 From: Weikai Xie <75515121+XieWeikai@users.noreply.github.com> Date: Thu, 20 Aug 2026 21:40:22 +0800 Subject: [PATCH] fix: respect gradient checkpointing configuration --- src/openpi/models_pytorch/gemma_pytorch.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/src/openpi/models_pytorch/gemma_pytorch.py b/src/openpi/models_pytorch/gemma_pytorch.py index 203b36be8a..d9a1c1dda3 100644 --- a/src/openpi/models_pytorch/gemma_pytorch.py +++ b/src/openpi/models_pytorch/gemma_pytorch.py @@ -134,13 +134,6 @@ def forward( and self.training ) or (hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training) - # Force enable gradient checkpointing if we're in training mode and the model supports it - if self.training and hasattr(self.gemma_expert.model, "gradient_checkpointing"): - if not self.gemma_expert.model.gradient_checkpointing: - print("Forcing gradient checkpointing to be enabled for Gemma expert model") - self.gemma_expert.model.gradient_checkpointing = True - use_gradient_checkpointing = True - # Debug gradient checkpointing status if hasattr(self, "_debug_gc_printed") and not self._debug_gc_printed: print(f"Gemma expert model gradient checkpointing: {use_gradient_checkpointing}")