Fix namedsharding

coolkp · coolkp · commit 6318fa6aac30 · 2025-06-25T19:34:13.000Z
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline.py b/src/maxdiffusion/pipelines/wan/wan_pipeline.py
@@ -397,7 +397,7 @@ def __call__(
             num_channels_latents=num_channel_latents,
         )
 
-      data_sharding = NamedSharding(self.devices_array, P())
+      data_sharding = NamedSharding(self.mesh, P())
       if len(prompt) % jax.device_count() == 0:
         data_sharding = jax.sharding.NamedSharding(self.mesh, P(*self.config.data_sharding))
 

Original file line number	Diff line number	Diff line change
`@@ -397,7 +397,7 @@ def __call__(`
`397`	`397`	`num_channels_latents=num_channel_latents,`
`398`	`398`	`)`
`399`	`399`
`400`		`- data_sharding = NamedSharding(self.devices_array, P())`
	`400`	`+ data_sharding = NamedSharding(self.mesh, P())`
`401`	`401`	`if len(prompt) % jax.device_count() == 0:`
`402`	`402`	`data_sharding = jax.sharding.NamedSharding(self.mesh, P(*self.config.data_sharding))`
`403`	`403`