test

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -285,8 +285,6 @@ def merge_lora_for_scanned(
 
   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
 
     nnx_path_str = ".".join(map(str, path))
     lora_key_template = translate_fn(nnx_path_str) if translate_fn else None
@@ -296,6 +294,13 @@ def merge_lora_for_scanned(
       max_logging.log(f"lora_key_template: {lora_key_template}")
       max_logging.log(f"Module: {module}")
       max_logging.log(f"Module Type: {type(module)}")
+      continue
+
+    if not isinstance(module, (nnx.Linear, nnx.Conv, nnx.LayerNorm, nnx.RMSNorm, nnx.Embed)):
+      continue
+
+    nnx_path_str = ".".join(map(str, path))
+    lora_key_template = translate_fn(nnx_path_str) if translate_fn else None
 
     if not lora_key_template:
       continue