Paged attn (#1036)

* nice code * device type adjustment Signed-off-by: Liu, Kaixuan <[email protected]>
huggingface · Nov 27, 2024 · 587837e · 587837e
1 parent 51030e5
commit 587837e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/optimum/exporters/ipex/modeling_utils.py b/optimum/exporters/ipex/modeling_utils.py
@@ -602,11 +602,11 @@ def __init__(self, module, config) -> None:
         self.q_slice = self.q_proj.out_features
         self.k_slice = self.q_slice + self.k_proj.out_features
         self.v_slice = self.k_slice + self.v_proj.out_features
-        if self.module_device == "cpu":
+        if self.module_device.type == "cpu":
             if module.o_proj.__class__.__name__ not in ["LinearAllreduce"]:
                 self.mha_linear_add = LinearAdd(module.o_proj)
 
-        elif self.module_device == "xpu":
+        elif self.module_device.type == "xpu":
             if module.o_proj.__class__.__name__ not in ["LinearAllreduce"]:
                 self.mha_linear_add = XPULinearAdd(module.o_proj)