From ce2691c8d331f5406958c62e424595193d6de390 Mon Sep 17 00:00:00 2001 From: ciaranbor Date: Tue, 16 Dec 2025 15:21:12 +0000 Subject: [PATCH] Fix shapes --- .../worker/engines/mflux/pipefusion/distributed_denoising.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py b/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py index ff938bcf..f26d7f5f 100644 --- a/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py +++ b/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py @@ -304,7 +304,7 @@ class DistributedDenoising: elif self.is_first_stage: hidden_states = mx.distributed.recv_like( - hidden_states, src=self.world_size - 1, group=self.group + prev_latents, src=self.world_size - 1, group=self.group ) mx.eval(hidden_states)