Enabling PEFT for jais-13b

When trying to fine-tune jais-13b using qlora, I encourted this error :
This error says that the "hidden_states" is leaf variable(Leaf Variable: A tensor that is not the result of an operation and has requires_grad=True.) therefore it doesn't accept in-place operation like in this error:
hidden_states *= torch.tensor(float(self.embeddings_scale), dtype=hidden_states.dtype, device=hidden_states.device )

modeling_jais.py

@@ -866,10 +866,13 @@ 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
-        )
-
+        
+        # hidden_states *= torch.tensor(
+        #     float(self.embeddings_scale), dtype=hidden_states.dtype, device=hidden_states.device
+        # )
+        aux_hidden = torch.tensor(float(self.embeddings_scale), dtype=hidden_states.dtype, device=hidden_states.device)
+        hidden_states = hidden_states * aux_hidden
+        
         if token_type_ids is not None:
             token_type_embeds = self.wte(token_type_ids)
             hidden_states = hidden_states + token_type_embeds