Fix dtype of unpacked tensor (#1840)

Kaihui-intel · web-flow · commit 29fdecbbb44c · 2024-06-05T09:24:49.000+08:00
Signed-off-by: Kaihui-intel &lt;kaihui.tang@intel.com&gt;
diff --git a/neural_compressor/torch/algorithms/weight_only/modules.py b/neural_compressor/torch/algorithms/weight_only/modules.py
@@ -289,7 +289,7 @@ def pack_tensor_with_torch(self, raw_tensor):
     def unpack_tensor_with_torch(self, packed_tensor):
         target_dtype = torch.int8 if not hasattr(self, "qzeros") or "int" not in self.dtype else torch.uint8
         target_len = packed_tensor.shape[1] * self.n_pack
-        unpacked_tensor = torch.zeros(packed_tensor.shape[0], target_len, dtype=self.compression_dtype).to(self.device)
+        unpacked_tensor = torch.zeros(packed_tensor.shape[0], target_len, dtype=target_dtype).to(self.device)
         mask = torch.tensor(2**self.bits - 1, dtype=self.compression_dtype).to(self.device)
         for j in range(packed_tensor.shape[1]):
             for e in range(self.n_pack):
