diff --git a/README.md b/README.md
index 058b6a5..12d41ee 100644
--- a/README.md
+++ b/README.md
@@ -11,98 +11,100 @@ The `ComfyUI-to-Python-Extension` is a powerful tool that translates ComfyUI wor
 
 ```
 import random
+import torch
 import sys
 
 sys.path.append("../")
 from nodes import (
+    VAEDecode,
     KSamplerAdvanced,
     EmptyLatentImage,
-    VAEDecodeTiled,
     SaveImage,
-    CLIPTextEncode,
     CheckpointLoaderSimple,
+    CLIPTextEncode,
 )
 
 
 def main():
-    checkpointloadersimple = CheckpointLoaderSimple()
-    checkpointloadersimple_4 = checkpointloadersimple.load_checkpoint(
-        ckpt_name="sd_xl_base_1.0.safetensors"
-    )
-
-    emptylatentimage = EmptyLatentImage()
-    emptylatentimage_5 = emptylatentimage.generate(
-        width=1024, height=1024, batch_size=1
-    )
-
-    cliptextencode = CLIPTextEncode()
-    cliptextencode_6 = cliptextencode.encode(
-        text="evening sunset scenery blue sky nature, glass bottle with a galaxy in it",
-        clip=checkpointloadersimple_4[1],
-    )
-
-    cliptextencode_7 = cliptextencode.encode(
-        text="text, watermark", clip=checkpointloadersimple_4[1]
-    )
-
-    checkpointloadersimple_12 = checkpointloadersimple.load_checkpoint(
-        ckpt_name="sd_xl_refiner_1.0.safetensors"
-    )
-
-    cliptextencode_15 = cliptextencode.encode(
-        text="evening sunset scenery blue sky nature, glass bottle with a galaxy in it",
-        clip=checkpointloadersimple_12[1],
-    )
-
-    cliptextencode_16 = cliptextencode.encode(
-        text="text, watermark", clip=checkpointloadersimple_12[1]
-    )
-
-    ksampleradvanced = KSamplerAdvanced()
-    vaedecodetiled = VAEDecodeTiled()
-    saveimage = SaveImage()
-
-    for q in range(10):
-        ksampleradvanced_10 = ksampleradvanced.sample(
-            add_noise="enable",
-            noise_seed=random.randint(1, 2**64),
-            steps=25,
-            cfg=8,
-            sampler_name="euler",
-            scheduler="normal",
-            start_at_step=0,
-            end_at_step=20,
-            return_with_leftover_noise="enable",
-            model=checkpointloadersimple_4[0],
-            positive=cliptextencode_6[0],
-            negative=cliptextencode_7[0],
-            latent_image=emptylatentimage_5[0],
+    with torch.inference_mode():
+        checkpointloadersimple = CheckpointLoaderSimple()
+        checkpointloadersimple_4 = checkpointloadersimple.load_checkpoint(
+            ckpt_name="sd_xl_base_1.0.safetensors"
         )
 
-        ksampleradvanced_11 = ksampleradvanced.sample(
-            add_noise="disable",
-            noise_seed=random.randint(1, 2**64),
-            steps=25,
-            cfg=8,
-            sampler_name="euler",
-            scheduler="normal",
-            start_at_step=20,
-            end_at_step=10000,
-            return_with_leftover_noise="disable",
-            model=checkpointloadersimple_12[0],
-            positive=cliptextencode_15[0],
-            negative=cliptextencode_16[0],
-            latent_image=ksampleradvanced_10[0],
+        emptylatentimage = EmptyLatentImage()
+        emptylatentimage_5 = emptylatentimage.generate(
+            width=1024, height=1024, batch_size=1
         )
 
-        vaedecodetiled_17 = vaedecodetiled.decode(
-            samples=ksampleradvanced_11[0], vae=checkpointloadersimple_12[2]
+        cliptextencode = CLIPTextEncode()
+        cliptextencode_6 = cliptextencode.encode(
+            text="evening sunset scenery blue sky nature, glass bottle with a galaxy in it",
+            clip=checkpointloadersimple_4[1],
         )
 
-        saveimage_19 = saveimage.save_images(
-            filename_prefix="ComfyUI", images=vaedecodetiled_17[0].detach()
+        cliptextencode_7 = cliptextencode.encode(
+            text="text, watermark", clip=checkpointloadersimple_4[1]
         )
 
+        checkpointloadersimple_12 = checkpointloadersimple.load_checkpoint(
+            ckpt_name="sd_xl_refiner_1.0.safetensors"
+        )
+
+        cliptextencode_15 = cliptextencode.encode(
+            text="evening sunset scenery blue sky nature, glass bottle with a galaxy in it",
+            clip=checkpointloadersimple_12[1],
+        )
+
+        cliptextencode_16 = cliptextencode.encode(
+            text="text, watermark", clip=checkpointloadersimple_12[1]
+        )
+
+        ksampleradvanced = KSamplerAdvanced()
+        vaedecode = VAEDecode()
+        saveimage = SaveImage()
+
+        for q in range(10):
+            ksampleradvanced_10 = ksampleradvanced.sample(
+                add_noise="enable",
+                noise_seed=random.randint(1, 2**64),
+                steps=25,
+                cfg=8,
+                sampler_name="euler",
+                scheduler="normal",
+                start_at_step=0,
+                end_at_step=20,
+                return_with_leftover_noise="enable",
+                model=checkpointloadersimple_4[0],
+                positive=cliptextencode_6[0],
+                negative=cliptextencode_7[0],
+                latent_image=emptylatentimage_5[0],
+            )
+
+            ksampleradvanced_11 = ksampleradvanced.sample(
+                add_noise="disable",
+                noise_seed=random.randint(1, 2**64),
+                steps=25,
+                cfg=8,
+                sampler_name="euler",
+                scheduler="normal",
+                start_at_step=20,
+                end_at_step=10000,
+                return_with_leftover_noise="disable",
+                model=checkpointloadersimple_12[0],
+                positive=cliptextencode_15[0],
+                negative=cliptextencode_16[0],
+                latent_image=ksampleradvanced_10[0],
+            )
+
+            vaedecode_17 = vaedecode.decode(
+                samples=ksampleradvanced_11[0], vae=checkpointloadersimple_12[2]
+            )
+
+            saveimage_19 = saveimage.save_images(
+                filename_prefix="ComfyUI", images=vaedecode_17[0].detach()
+            )
+
 
 if __name__ == "__main__":
     main()
diff --git a/comfyui_to_python.py b/comfyui_to_python.py
index c31d7b9..7e79674 100644
--- a/comfyui_to_python.py
+++ b/comfyui_to_python.py
@@ -197,10 +197,6 @@ def get_class_info(class_type: str) -> (str, str, str):
         import_statement (str): Import statement string
         class_code (str): Class initialization code
     """
-    # If the class is 'VAEDecode', adjust the class name
-    if class_type == 'VAEDecode':
-        class_type = 'VAEDecodeTiled'
-    
     import_statement = class_type
     class_code = f'{class_type.lower()} = {class_type}()'
 
@@ -219,13 +215,12 @@ def assemble_python_code(import_statements: set, loader_code: List[str], code: L
     Returns:
         final_code (str): Generated final code as a string
     """
-    static_imports = ['import random']
+    static_imports = ['import random', 'import torch']
     imports_code = [f"from nodes import {', '.join([class_name for class_name in import_statements])}" ]
-    main_function_code = f"def main():\n\t" + '\n\t'.join(loader_code) + f'\n\n\tfor q in range({queue_size}):\n\t' + '\n\t'.join(code)
+    main_function_code = f"def main():\n\t" + 'with torch.inference_mode():\n\t\t' + '\n\t\t'.join(loader_code) + f'\n\n\t\tfor q in range({queue_size}):\n\t\t' + '\n\t\t'.join(code)
     final_code = '\n'.join(static_imports + ['import sys\nsys.path.append("../")'] + imports_code + ['', main_function_code, '', 'if __name__ == "__main__":', '\tmain()'])
     final_code = black.format_str(final_code, mode=black.Mode())
 
-
     return final_code