Update merged_files3.py

merged_files3.py

@@ -492,11 +492,10 @@ device = torch.device('cuda')
 dtype = torch.float16  # Use float16 consistently for all models
 
 
-sd_offset = sf.load_file(model_path, device=device) # Use device variable
+sd_offset = sf.load_file(model_path) # Use device variable
 sd_origin = unet.state_dict()
 keys = sd_origin.keys()
-sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()}
-unet.load_state_dict(sd_merged, strict=True)
+sd_offset = {k: v.to(device) for k, v in sd_offset.items()} # Move each tensor to GPUunet.load_state_dict(sd_merged, strict=True)
 del sd_offset, sd_origin, sd_merged, keys