fix: pass all arguments through gradient checkpoint in BasicTransformerBlock

[Mr-Neutr0n]

Summary

BasicTransformerBlock.forward() silently drops additional_tokens and n_times_crossframe_attn_in_self when gradient checkpointing is enabled. The checkpoint() call only forwards x and context, causing the other two arguments to revert to their defaults (None and 0):

# Before (bug)
return checkpoint(self._forward, x, context)

This means that during training with checkpointing enabled:

• additional_tokens is silently ignored (defaults to None), so any extra conditioning tokens passed to self-attention and cross-attention are lost.
• n_times_crossframe_attn_in_self is silently ignored (defaults to 0), so cross-frame attention repeat behavior in self-attention is disabled.

The non-checkpointed path (self._forward(**kwargs)) correctly passes all arguments, so this bug only manifests when checkpoint=True — the default setting — making it easy to miss.

Fix

Pass all four arguments through checkpoint():

# After (fix)
return checkpoint(self._forward, x, context, additional_tokens, n_times_crossframe_attn_in_self)

This ensures _forward receives the same arguments regardless of whether gradient checkpointing is enabled.