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

Author: pollockjj

distorch_2.py

@@ -301,6 +301,8 @@ def new_partially_load(self, device_to, extra_memory=0, full_load=False, force_p
 
             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")