scale conv2d as well

cloneofsimo · Jan 14, 2023 · 204d6ed · 204d6ed
1 parent 7dca434
commit 204d6ed
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lora_diffusion/lora.py b/lora_diffusion/lora.py
@@ -749,7 +749,7 @@ def monkeypatch_add_lora(
 
 def tune_lora_scale(model, alpha: float = 1.0):
     for _module in model.modules():
-        if _module.__class__.__name__ == "LoraInjectedLinear":
+        if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d"]:
             _module.scale = alpha