diff --git a/examples/wan2.1_fun/predict_i2v.py b/examples/wan2.1_fun/predict_i2v.py
index 5c1e3de..96a2f0d 100755
--- a/examples/wan2.1_fun/predict_i2v.py
+++ b/examples/wan2.1_fun/predict_i2v.py
@@ -246,7 +246,7 @@ if cfg_skip_ratio is not None:
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    pipeline = merge_lora(pipeline, lora_path, lora_weight)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
 
 with torch.no_grad():
     video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
@@ -274,7 +274,7 @@ with torch.no_grad():
     ).videos
 
 if lora_path is not None:
-    pipeline = unmerge_lora(pipeline, lora_path, lora_weight)
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
 
 def save_results():
     if not os.path.exists(save_path):
diff --git a/examples/wan2.1_fun/predict_t2v.py b/examples/wan2.1_fun/predict_t2v.py
index 124fbc1..e49dac3 100755
--- a/examples/wan2.1_fun/predict_t2v.py
+++ b/examples/wan2.1_fun/predict_t2v.py
@@ -253,7 +253,7 @@ if cfg_skip_ratio is not None:
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    pipeline = merge_lora(pipeline, lora_path, lora_weight)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
 
 with torch.no_grad():
     video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
@@ -293,7 +293,7 @@ with torch.no_grad():
         ).videos
 
 if lora_path is not None:
-    pipeline = unmerge_lora(pipeline, lora_path, lora_weight)
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
 
 def save_results():
     if not os.path.exists(save_path):