src/transformers/modeling_flax_utils.py/FlaxPreTrainedModel/enable_gradient_checkpointing
class FlaxPreTrainedModel:    def enable_gradient_checkpointing(self):
        raise NotImplementedError(f"gradient checkpointing method has to be implemented for {self}")