Change in wan_vae_test.py

Author: prishajain1

src/maxdiffusion/tests/wan_vae_test.py

@@ -266,7 +266,7 @@ def test_wan_resample(self):
       # channels is always last here
       input_shape = (batch, t, h, w, dim)
       dummy_input = jnp.ones(input_shape)
-      output = wan_resample(dummy_input)
+      output, _, _ = wan_resample(dummy_input)
       assert output.shape == (batch, t, h // 2, w // 2, dim)
 
   def test_3d_conv(self):
@@ -347,7 +347,7 @@ def test_wan_residual(self):
     with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
       wan_residual_block = WanResidualBlock(in_dim=in_dim, out_dim=out_dim, rngs=rngs, mesh=mesh)
       dummy_input = jnp.ones(input_shape)
-      dummy_output = wan_residual_block(dummy_input)
+      dummy_output, _, _ = wan_residual_block(dummy_input)
       assert dummy_output.shape == expected_output_shape
       # --- Test Case 1: different in/out dim ---
       in_dim = 96
@@ -356,7 +356,7 @@ def test_wan_residual(self):
 
       wan_residual_block = WanResidualBlock(in_dim=in_dim, out_dim=out_dim, rngs=rngs, mesh=mesh)
       dummy_input = jnp.ones(input_shape)
-      dummy_output = wan_residual_block(dummy_input)
+      dummy_output, _, _ = wan_residual_block(dummy_input)
       assert dummy_output.shape == expected_output_shape
 
   def test_wan_attention(self):
@@ -371,7 +371,7 @@ def test_wan_attention(self):
     with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
       wan_attention = WanAttentionBlock(dim=dim, rngs=rngs)
       dummy_input = jnp.ones(input_shape)
-      output = wan_attention(dummy_input)
+      output, _, _ = wan_attention(dummy_input)
       assert output.shape == input_shape
 
   def test_wan_midblock(self):
@@ -396,7 +396,7 @@ def test_wan_midblock(self):
     with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
       wan_midblock = WanMidBlock(dim=dim, rngs=rngs, mesh=mesh)
       dummy_input = jnp.ones(input_shape)
-      output = wan_midblock(dummy_input)
+      output, _, _ = wan_midblock(dummy_input)
       assert output.shape == input_shape
 
   def test_wan_decode(self):