src/transformers/models/longt5/modeling_flax_longt5.py/FlaxLongT5PreTrainedModel/enable_gradient_checkpointing
class FlaxLongT5PreTrainedModel:    def enable_gradient_checkpointing(self):
        self._module = self.module_class(
            config=self.config,
            dtype=self.dtype,
            gradient_checkpointing=True,
        )