Fix

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -224,7 +224,9 @@ def merge_lora(model: nnx.Module, state_dict: dict, scale: float, translate_fn=N
 
     assigned_count = 0
     for path, module in nnx.iter_graph(model):
-        if not isinstance(module, (nnx.Linear, nnx.Conv, nnx.LayerNorm, nnx.RMSNorm, nnx.Embed)): continue
+        if not isinstance(module, (nnx.Linear, nnx.Conv, nnx.LayerNorm, nnx.RMSNorm, nnx.Embed)): 
+          max_logging.log(f"Skipping non-supported module type: {module}")
+          continue
 
         nnx_path_str = ".".join(map(str, path))
         lora_key = translate_fn(nnx_path_str) if translate_fn else None
@@ -319,7 +321,9 @@ def merge_lora_for_scanned(model: nnx.Module, state_dict: dict, scale: float, tr
 
     assigned_count = 0
     for path, module in nnx.iter_graph(model):
-        if not isinstance(module, (nnx.Linear, nnx.Conv, nnx.LayerNorm, nnx.RMSNorm, nnx.Embed)): continue
+        if not isinstance(module, (nnx.Linear, nnx.Conv, nnx.LayerNorm, nnx.RMSNorm, nnx.Embed)):
+          max_logging.log(f"Skipping non-supported module type: {module}")
+          continue
 
         nnx_path_str = ".".join(map(str, path))
         lora_key_template = translate_fn(nnx_path_str) if translate_fn else None