Fix shapes

2026-01-06 10:51:21 +00:00
parent 076d2901e8
commit ce2691c8d3
1 changed files with 1 additions and 1 deletions
@@ -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)