diff --git a/modules/impact/config.py b/modules/impact/config.py
index 7f7c365c..d01ae427 100644
--- a/modules/impact/config.py
+++ b/modules/impact/config.py
@@ -1,7 +1,7 @@
 import configparser
 import os
 
-version_code = [6, 0, 1]
+version_code = [6, 0, 2]
 version = f"V{version_code[0]}.{version_code[1]}" + (f'.{version_code[2]}' if len(version_code) > 2 else '')
 
 dependency_version = 22
diff --git a/modules/impact/impact_sampling.py b/modules/impact/impact_sampling.py
index 151e6b14..5c74a8d7 100644
--- a/modules/impact/impact_sampling.py
+++ b/modules/impact/impact_sampling.py
@@ -6,6 +6,7 @@
 import comfy
 import torch
 import math
+import comfy.model_management as mm
 
 
 try:
@@ -151,6 +152,13 @@ def touched_callback(step, x0, x, total_steps):
         positive = node_helpers.conditioning_set_values(positive, {"guidance": cfg})
         guider.set_conds(positive)
 
+    device = mm.get_torch_device()
+
+    noise = noise.to(device)
+    latent_image = latent_image.to(device)
+    if noise_mask is not None:
+        noise_mask = noise_mask.to(device)
+
     samples = guider.sample(noise, latent_image, sampler, sigmas, denoise_mask=noise_mask, callback=touched_callback, disable_pbar=disable_pbar, seed=noise_seed)
     samples = samples.to(comfy.model_management.intermediate_device())
 
diff --git a/pyproject.toml b/pyproject.toml
index 18531826..b61afb70 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-impact-pack"
 description = "This extension offers various detector nodes and detailer nodes that allow you to configure a workflow that automatically enhances facial details. And provide iterative upscaler."
-version = "6.0.1"
+version = "6.0.2"
 license = { file = "LICENSE.txt" }
 dependencies = ["segment-anything", "scikit-image", "piexif", "transformers", "opencv-python-headless", "GitPython", "scipy>=1.11.4"]