diff --git a/comfy_extras/nodes_autoregressive.py b/comfy_extras/nodes_autoregressive.py
index 39c5b9952..6dfcaf766 100644
--- a/comfy_extras/nodes_autoregressive.py
+++ b/comfy_extras/nodes_autoregressive.py
@@ -18,7 +18,7 @@ class AutoRegressiveGeneration:
             }
         }
 
-    RETURN_TYPES = ("TOKENS",)
+    RETURN_TYPES = ("GENERATED_TOKENS",)
     FUNCTION = "generate"
 
     CATEGORY = "sampling"
@@ -35,7 +35,7 @@ class DecodeTokens:
         return {
             "required": {
                 "clip": (IO.CLIP, {"tooltip": "The model used for generation."}),
-                "tokens": ("TOKENS", ),}
+                "tokens": ("GENERATED_TOKENS", ),}
         }
 
     FUNCTION = "decode"