While fine-tuning DreamZero-AgiBot on a bimanual YAM dataset we wanted a deterministic validation loss, i.e. run the exact training forward on held-out batches with gradients disabled. That crashes inside the DiT:
AttributeError: 'tuple' object has no attribute 'shape'
Root cause: in wan_video_dit_action_casual_chunk.py (CausalWanModel, ~line 2143), the block loop only unpacks the (x, kv_cache) return in the gradient-checkpointing branch. The non-checkpointed branch does x = block(x, **kwargs), so as soon as torch.is_grad_enabled() is false the next block receives a tuple.
Normal training never hits this because checkpointing is on and grads are enabled, so the bug is latent, but it breaks any eval-style use of the training objective (validation loss curves, offline checkpoint selection, loss-parity checks after checkpoint surgery).
One-line fix incoming. We ran with it for a 12k-step fine-tune and used the grad-disabled forward to score ~26 milestone checkpoints; the run and eval code are at https://github.com/robocurve/dreamzero-yam if useful.
While fine-tuning DreamZero-AgiBot on a bimanual YAM dataset we wanted a deterministic validation loss, i.e. run the exact training forward on held-out batches with gradients disabled. That crashes inside the DiT:
Root cause: in
wan_video_dit_action_casual_chunk.py(CausalWanModel, ~line 2143), the block loop only unpacks the(x, kv_cache)return in the gradient-checkpointing branch. The non-checkpointed branch doesx = block(x, **kwargs), so as soon astorch.is_grad_enabled()is false the next block receives a tuple.Normal training never hits this because checkpointing is on and grads are enabled, so the bug is latent, but it breaks any eval-style use of the training objective (validation loss curves, offline checkpoint selection, loss-parity checks after checkpoint surgery).
One-line fix incoming. We ran with it for a 12k-step fine-tune and used the grad-disabled forward to score ~26 milestone checkpoints; the run and eval code are at https://github.com/robocurve/dreamzero-yam if useful.