diff --git a/efficiency_nodes.py b/efficiency_nodes.py
index f4f7040..b7806da 100644
--- a/efficiency_nodes.py
+++ b/efficiency_nodes.py
@@ -1410,7 +1410,7 @@ class TSC_KSampler:
                     capsule_result = xy_capsule.get_result(model, clip, vae)
                     if capsule_result is not None:
                         image, latent = capsule_result
-                        latent_list.append(latent['samples'])
+                        latent_list.append(latent)
 
                 if capsule_result is None:
                     if preview_method != "none":
@@ -1428,6 +1428,9 @@ class TSC_KSampler:
                     # Decode the latent tensor
                     image = vae_decode_latent(vae, samples, vae_decode)
 
+                    if xy_capsule is not None:
+                        xy_capsule.set_result(image, samples)
+
                 # Add the resulting image tensor to image_tensor_list
                 image_tensor_list.append(image)