Update modeling_jais.py

modeling_jais.py

@@ -535,6 +535,11 @@ 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
@@ -861,6 +866,7 @@ class JAISModel(JAISPreTrainedModel):
             hidden_states = inputs_embeds + position_embeds
         else:
             hidden_states = inputs_embeds
+        '''hidden_states *= torch.tensor(float(self.embeddings_scale), dtype=hidden_states.dtype, device=hidden_states.device) '''  
         scale_factor_hidden = torch.tensor(float(self.embeddings_scale), dtype=hidden_states.dtype, device=hidden_states.device)
         hidden_states = hidden_states * scale_factor_hidden