Fix issue with layer offloading on ernie

extensions_built_in/diffusion_models/ernie_image/ernie_image.py

@@ -110,8 +110,7 @@ class ErnieImageModel(BaseModel):
                 self.device_torch,
                 offload_percent=self.model_config.layer_offloading_transformer_percent,
                 ignore_modules=[
-                    transformer.x_pad_token,
-                    transformer.cap_pad_token,
+                    transformer.x_embedder,
                 ],
             )
 

extensions_built_in/diffusion_models/ernie_image/transformer.py

@@ -338,6 +338,12 @@ class ErnieImageTransformer2DModel(ModelMixin, ConfigMixin):
         nn.init.zeros_(self.final_linear.weight)
         nn.init.zeros_(self.final_linear.bias)
         self.gradient_checkpointing = False
+        self.onload_device = None
+    
+    @property
+    def device(self):
+        # use self.x_embeddersince we ignore it in memory management
+        return next(self.x_embedder.parameters()).device
 
     def forward(
         self,