fix

prishajain1 · prishajain1 · commit cdc0271d8d32 · 2026-03-23T21:42:17.000+05:30
diff --git a/src/maxdiffusion/tests/wan_magcache_test.py b/src/maxdiffusion/tests/wan_magcache_test.py
@@ -252,7 +252,7 @@ def setUpClass(cls):
     )
     cls.config = pyconfig.config
     checkpoint_loader = WanCheckpointerI2V_2_1(cls.config)
-    cls.pipeline, _ = checkpoint_loader.load_checkpoint()
+    cls.pipeline, _, _ = checkpoint_loader.load_checkpoint()
 
     cls.image = load_image(cls.config.image_url)
     cls.prompt = [cls.config.prompt] * cls.config.global_batch_size_to_train_on

Original file line number	Diff line number	Diff line change
`@@ -252,7 +252,7 @@ def setUpClass(cls):`
`252`	`252`	`)`
`253`	`253`	`cls.config = pyconfig.config`
`254`	`254`	`checkpoint_loader = WanCheckpointerI2V_2_1(cls.config)`
`255`		`- cls.pipeline, _ = checkpoint_loader.load_checkpoint()`
	`255`	`+ cls.pipeline, _, _ = checkpoint_loader.load_checkpoint()`
`256`	`256`
`257`	`257`	`cls.image = load_image(cls.config.image_url)`
`258`	`258`	`cls.prompt = [cls.config.prompt] * cls.config.global_batch_size_to_train_on`