use_real correction

prishajain1 · prishajain1 · commit 8cdf7718c691 · 2025-12-17T11:20:18.000+05:30
diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan.py b/src/maxdiffusion/models/wan/transformers/transformer_wan.py
@@ -515,7 +515,6 @@ def init_block(rngs):
           precision=precision,
           attention=attention,
           dropout=dropout,
-          use_real=False,
       )
 
     self.gradient_checkpoint = GradientCheckpointType.from_str(remat_policy)
@@ -544,7 +543,6 @@ def init_block(rngs):
             precision=precision,
             attention=attention,
             enable_jax_named_scopes=enable_jax_named_scopes,
-            use_real=False,
         )
         blocks.append(block)
       self.blocks = blocks

Original file line number	Diff line number	Diff line change
`@@ -515,7 +515,6 @@ def init_block(rngs):`
`515`	`515`	`precision=precision,`
`516`	`516`	`attention=attention,`
`517`	`517`	`dropout=dropout,`
`518`		`- use_real=False,`
`519`	`518`	`)`
`520`	`519`
`521`	`520`	`self.gradient_checkpoint = GradientCheckpointType.from_str(remat_policy)`
`@@ -544,7 +543,6 @@ def init_block(rngs):`
`544`	`543`	`precision=precision,`
`545`	`544`	`attention=attention,`
`546`	`545`	`enable_jax_named_scopes=enable_jax_named_scopes,`
`547`		`- use_real=False,`
`548`	`546`	`)`
`549`	`547`	`blocks.append(block)`
`550`	`548`	`self.blocks = blocks`