ensure correct device placement

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

Author: kylesayrs

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -47,6 +47,7 @@
     iter_named_leaf_modules,
 )
 from compressed_tensors.utils import (
+    get_execution_device,
     get_safetensors_folder,
     has_offloaded_params,
     merge_names,
@@ -403,12 +404,14 @@ def compress_model(self, model: Module):
                     )
 
                 # remove any existing parameters
+                device = get_execution_device(module)
                 for name, _ in list(module.named_parameters()):
                     delattr(module, name)
 
                 # replace with compressed parameters
                 for name, value in state_dict.items():
                     name = name.removeprefix(f"{prefix}.")
+                    value = value.to(device)
                     param = torch.nn.Parameter(value, requires_grad=False)
                     register_offload_parameter(module, name, param)
 
@@ -456,12 +459,14 @@ def decompress_model(self, model: Module):
                     }
 
                 # remove any existing parameters
+                device = get_execution_device(module)
                 for name, _ in list(module.named_parameters()):
                     delattr(module, name)
 
                 # replace with decompressed parameters
                 for name, value in state_dict.items():
                     name = name.removeprefix(f"{prefix}.")
+                    value = value.to(device)
                     param = torch.nn.Parameter(value, requires_grad=False)
                     register_offload_parameter(module, name, param)