src/transformers/models/electra/modeling_flax_electra.py/FlaxElectraPreTrainedModel/enable_gradient_checkpointing
class FlaxElectraPreTrainedModel:    def enable_gradient_checkpointing(self):
        self._module = self.module_class(
            config=self.config,
            dtype=self.dtype,
            gradient_checkpointing=True,
        )