fix

prishajain1 · prishajain1 · commit 885b4415de8a · 2026-02-27T22:38:08.000+05:30
diff --git a/src/maxdiffusion/models/ltx2/attention_ltx2.py b/src/maxdiffusion/models/ltx2/attention_ltx2.py
@@ -388,6 +388,8 @@ def __init__(
         heads=heads,
         dim_head=dim_head,
         dtype=dtype,
+        axis_names_q=("batch", "heads", "length", "kv"),
+        axis_names_kv=("batch", "heads", "length", "kv"),
     )
 
   def __call__(

Original file line number	Diff line number	Diff line change
`@@ -388,6 +388,8 @@ def __init__(`
`388`	`388`	`heads=heads,`
`389`	`389`	`dim_head=dim_head,`
`390`	`390`	`dtype=dtype,`
	`391`	`+ axis_names_q=("batch", "heads", "length", "kv"),`
	`392`	`+ axis_names_kv=("batch", "heads", "length", "kv"),`
`391`	`393`	`)`
`392`	`394`
`393`	`395`	`def __call__(`