Fix

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -17,6 +17,7 @@
 from typing import Union, Tuple, Optional
 import re
 import torch
+import jax
 from jax import dlpack
 import jax.numpy as jnp
 from flax import nnx
@@ -155,8 +156,8 @@ def __call__(self, x):
 
 def _to_jax_array(v):
     if isinstance(v, torch.Tensor):
-        return dlpack.from_dlpack(v)
-    return jnp.array(v)
+        return jax.device_put(dlpack.from_dlpack(v))
+    return jax.device_put(jnp.array(v))
 
 def merge_lora(model: nnx.Module, state_dict: dict, scale: float, translate_fn=None):
     """