Update modeling_ernie4_5_moe.py

modeling_ernie4_5_moe.py

@@ -1113,7 +1113,11 @@ class Ernie4_5_Model(Ernie4_5_PretrainedModel):
             past_key_values = DynamicCache()
 
         if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
+            # Account for CPU offloaded embed_tokens
+            embed_device = self.embed_tokens.weight.device
+            inputs_embeds = self.embed_tokens(input_ids.to(embed_device, non_blocking = True)).to(input_ids.device)
+        if not self.training:
+            inputs_embeds.requires_grad_(False)
         
         inputs_embeds = inputs_embeds.to(self.embed_tokens.weight.dtype)