maintain module_device (#384)

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -392,8 +392,8 @@ def compress_model(self, model: Module):
         for prefix, module in tqdm(model.named_modules(), desc="Compressing model"):
 
             if prefix in module_to_scheme or prefix in sparse_compression_targets:
-                module_device = get_execution_device(module).type
-                is_meta = (module_device == "meta")
+                module_device = get_execution_device(module)
+                is_meta = module_device.type == "meta"
 
                 exec_device = "meta" if is_meta else "cpu"
                 onloading_device = "meta" if is_meta else module_device

src/compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py

@@ -178,9 +178,13 @@ def sparse24_bitmask_compress(
 
     if tensor.is_meta:
         num_rows, num_cols = tensor.shape
-        compressed_values = torch.empty((num_rows, num_cols // 2), dtype=tensor.dtype, device="meta")
+        compressed_values = torch.empty(
+            (num_rows, num_cols // 2), dtype=tensor.dtype, device="meta"
+        )
         packed_cols = (num_cols + 7) // 8
-        bitmasks_packed = torch.empty((num_rows, packed_cols), dtype=torch.uint8, device="meta")
+        bitmasks_packed = torch.empty(
+            (num_rows, packed_cols), dtype=torch.uint8, device="meta"
+        )
         return compressed_values, bitmasks_packed
 
     bytemasks = get_24_bytemasks(tensor=tensor)

tests/test_compressors/model_compressors/test_model_compressor.py

@@ -446,10 +446,7 @@ def test_compress_model_meta(model_stub, q_format, s_config):
         cpu_model, s_config, q_format
     )
     # Only stores dtype because meta model does not store values
-    expected = {
-        k: v.dtype
-        for k, v in reference_compressor.compress(cpu_model).items()
-    }
+    expected = {k: v.dtype for k, v in reference_compressor.compress(cpu_model).items()}
 
     # Load model on meta device
     meta_model = AutoModelForCausalLM.from_pretrained(