From dd84cc9ca2c2fd0dd5d15f084e5bc2e7c31edf4e Mon Sep 17 00:00:00 2001
From: ciaranbor <ciaranborourke-dev@proton.me>
Date: Mon, 15 Dec 2025 16:57:02 +0000
Subject: [PATCH] Fix return

---
 .../worker/engines/mflux/pipefusion/distributed_denoising.py    | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py b/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py
index bac8bff2..d39c3b06 100644
--- a/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py
+++ b/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py
@@ -319,6 +319,8 @@ class DistributedDenoising:
             -hidden_states.shape[0] :
         ]
 
+        return hidden_states
+
     def _async_pipeline(
         self,
         t: int,