fix

Author: prishajain1

src/maxdiffusion/tests/test_ltx2_utils.py

@@ -185,7 +185,7 @@ def test_load_vocoder_weights(self):
             else:
                 pt_shape = shape
 
-            pt_weights[pt_key] = torch.randn(pt_shape)
+            pt_weights[pt_key] = jnp.array(torch.randn(pt_shape).numpy())
 
         with mock.patch("maxdiffusion.models.ltx2.ltx2_utils.load_sharded_checkpoint", return_value=pt_weights):
             loaded_weights = load_vocoder_weights(