diff --git a/dynthres_comfyui.py b/dynthres_comfyui.py
index 23112bb..4f17f39 100644
--- a/dynthres_comfyui.py
+++ b/dynthres_comfyui.py
@@ -35,6 +35,7 @@ class DynamicThresholdingComfyNode:
             uncond = input - args["uncond"]
             cond_scale = args["cond_scale"]
             time_step = model.model.model_sampling.timestep(args["sigma"])
+            time_step = time_step[0].item()
             dynamic_thresh.step = 999 - time_step
 
             return input - dynamic_thresh.dynthresh(cond, uncond, cond_scale, None)
@@ -69,6 +70,7 @@ class DynamicThresholdingSimpleComfyNode:
             uncond = input - args["uncond"]
             cond_scale = args["cond_scale"]
             time_step = model.model.model_sampling.timestep(args["sigma"])
+            time_step = time_step[0].value()
             dynamic_thresh.step = 999 - time_step
 
             return input - dynamic_thresh.dynthresh(cond, uncond, cond_scale, None)