From a3e807ff0dd1bd77cdb9dd85262b49fb442ac4e0 Mon Sep 17 00:00:00 2001 From: Jay Chooi Date: Thu, 6 Aug 2026 22:07:27 -0700 Subject: [PATCH] Unpack block output tuple in the non-checkpointed training branch Transformer blocks return (x, kv_cache). Only the gradient-checkpointing branch unpacked the tuple, so running the training forward under torch.no_grad() handed a tuple to the next block and crashed. Unpack it in the non-checkpointed branch too. --- .../dreamzero/modules/wan_video_dit_action_casual_chunk.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py b/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py index c5e25284..1f45c2ed 100644 --- a/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py +++ b/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py @@ -2140,7 +2140,12 @@ def custom_forward(*inputs, **kwargs): use_reentrant=False, ) else: - x = block(x, **kwargs) + # Blocks return (x, kv_cache); only the checkpointed branch + # unpacked it, so any grad-disabled training forward (e.g. + # deterministic eval of the training objective) crashed with + # AttributeError: 'tuple' object has no attribute 'shape'. + x, _updated_kv_cache = block(x, **kwargs) + assert _updated_kv_cache is None if clean_x is not None: x = x[:, clean_x.shape[1]:]