Update modeling_jais.py

modeling_jais.py

@@ -535,11 +535,6 @@ class JAISPreTrainedModel(PreTrainedModel):
                 stddev = self.config.initializer_range * mup_init_scale / math.sqrt(2 * self.config.n_layer)
                 p.data.normal_(mean=0.0, std=stddev)
 
-    def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, JAISModel):
-            module.gradient_checkpointing = value
-
-
 JAIS_START_DOCSTRING = r"""
 
     This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the