Parametrize test

coolkp · coolkp · commit 1c36022e50b2 · 2025-11-13T18:37:30.000Z
Signed-off-by: Kunjan Patel &lt;kunjanp@google.com&gt;
diff --git a/src/maxdiffusion/tests/wan_transformer_test.py b/src/maxdiffusion/tests/wan_transformer_test.py
@@ -188,7 +188,7 @@ def test_wan_block(self):
         qk_norm=qk_norm,
         cross_attn_norm=cross_attn_norm,
         eps=eps,
-        attention=attention,
+        attention="flash",
         mesh=mesh,
         flash_block_sizes=flash_block_sizes,
     )

Original file line number	Diff line number	Diff line change
`@@ -188,7 +188,7 @@ def test_wan_block(self):`
`188`	`188`	`qk_norm=qk_norm,`
`189`	`189`	`cross_attn_norm=cross_attn_norm,`
`190`	`190`	`eps=eps,`
`191`		`- attention=attention,`
	`191`	`+ attention="flash",`
`192`	`192`	`mesh=mesh,`
`193`	`193`	`flash_block_sizes=flash_block_sizes,`
`194`	`194`	`)`