[fix]: match Cosmos3 oracle attention contract

2026-07-21 16:05:37 -07:00
parent ea7c8939f5
commit 71ef62f336
1 changed files with 1 additions and 1 deletions
@@ -392,7 +392,7 @@ def _official_layer_parity(
            torch.cat([value_und, value_gen], dim=0),
            is_causal=False,
        )
-        return from_und_gen_splits(und_output, gen_output, query_pack)
+        return from_und_gen_splits(und_output, gen_output, query_pack), None

    block.self_attn.dispatch_attention_fn = exact_dispatch
    und, gen, cos_und, sin_und, cos_gen, sin_gen = (