From c713f637b0aa7823a1254a216be173ffaacde03e Mon Sep 17 00:00:00 2001
From: John Pollock <pollockjj@gmail.com>
Date: Wed, 15 Oct 2025 21:37:33 -0500
Subject: [PATCH] fix: ensure model.device is set in
 ModelPatcher.partially_load

Assign self.model.device = device_to during DisTorch V2 partially_load so the model's device reflects the target allocation after loading.

AssertionError: Input tensors must be on cuda.
Fixes #119

Possible issue when used with custom samplers
Fixes #130
---
 distorch_2.py | 2 ++
 1 file changed, 2 insertions(+)

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")