diff --git a/distorch_2.py b/distorch_2.py
index 1922875..fea6181 100644
--- a/distorch_2.py
+++ b/distorch_2.py
@@ -301,6 +301,8 @@ def register_patched_safetensor_modelpatcher():
 
             self.model.current_weight_patches_uuid = self.patches_uuid
 
+            self.model.device = device_to
+
             logger.info("[MultiGPU DisTorch V2] DisTorch loading completed.")
             logger.info(f"[MultiGPU DisTorch V2] Total memory: {mem_counter / (1024 * 1024):.2f}MB")