Remove extra params loading error

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -160,6 +160,8 @@ def create_model(rngs: nnx.Rngs, wan_config: dict):
   for path, val in flax.traverse_util.flatten_dict(params).items():
     if restored_checkpoint:
       path = path[:-1]
+    if path not in logical_state_sharding:
+       continue
     sharding = logical_state_sharding[path].value
     state[path].value = device_put_replicated(val, sharding)
   state = nnx.from_flat_state(state)