transformwe weight loading fix

prishajain1 · prishajain1 · commit 17c2903570fb · 2026-04-12T18:26:53.000+05:30
diff --git a/src/maxdiffusion/models/ltx2/attention_ltx2.py b/src/maxdiffusion/models/ltx2/attention_ltx2.py
@@ -437,8 +437,6 @@ def __init__(
           rngs=rngs,
           dtype=dtype,
       )
-    else:
-      self.to_gate_logits = None
 
     self.attention_op = NNXAttentionOp(
         mesh=mesh,

Original file line number	Diff line number	Diff line change
`@@ -437,8 +437,6 @@ def __init__(`
`437`	`437`	`rngs=rngs,`
`438`	`438`	`dtype=dtype,`
`439`	`439`	`)`
`440`		`- else:`
`441`		`- self.to_gate_logits = None`
`442`	`440`
`443`	`441`	`self.attention_op = NNXAttentionOp(`
`444`	`442`	`mesh=mesh,`