Fix error in lora_layers causing lora training to fail

2025-06-12 23:36:00 +02:00 · 2025-06-12 23:36:00 +02:00 · d2296ec197
commit d2296ec197
parent 4eb3f90fc9
1 changed files with 2 additions and 1 deletions
--- a/src/mflux/dreambooth/lora_layers/lora_layers.py
+++ b/src/mflux/dreambooth/lora_layers/lora_layers.py
@ -246,7 +246,8 @@ class LoRALayers:
        if block[key].get(name, False) and val.get(name, False):
            if name == "to_out" and isinstance(block[key][name], list):
                if len(block[key][name]) > 0:
-                    block[key][name][0] = val[name]
+                    lora_layer = val[name][0] if isinstance(val[name], list) else val[name]
+                    block[key][name][0] = lora_layer
            else:
                block[key][name] = val[name]