diff --git a/mat_anyone2.py b/mat_anyone2.py
index 7cfe8b3..00e9ffb 100644
--- a/mat_anyone2.py
+++ b/mat_anyone2.py
@@ -97,7 +97,7 @@ def inference_matanyone2(
             mask_t, processor, n_warmup, repeated_frames, pbar, 0, total_len
         )
         mask_out = processor.output_prob_to_mask(output_prob)
-        phas[0] = mask_out.unsqueeze(0).unsqueeze(0)
+        phas[0] = mask_out.unsqueeze(0).unsqueeze(0).cpu()
 
         for ti in range(1, length):
             image = vframes[ti].to(device)
diff --git a/pyproject.toml b/pyproject.toml
index 609e6e3..2e8d26f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "ComfyUI-MatAnyone"
 description = "MatAnyone in ComfyUI (Remove background)"
-version = "2.0"
+version = "2.1"
 license = {file = "LICENSE"}
 dependencies = ["torch", "omegaconf"]