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,