WIP to add the caption_proj weight to pixart sigma TE adapter

2024-07-06 13:00:21 -06:00
parent acb06d6ff3
commit cab8a1c7b8
8 changed files with 500 additions and 23 deletions
--- a/toolkit/network_mixins.py
+++ b/toolkit/network_mixins.py
@@ -289,7 +289,13 @@ class ToolkitModuleMixin:
            scaled_lora_weight = lora_weight * scale
            scaled_lora_output = scaled_lora_output + self.apply_dora(lx, scaled_lora_weight)

-        x = org_forwarded + scaled_lora_output
+        try:
+            x = org_forwarded + scaled_lora_output
+        except RuntimeError as e:
+            print(e)
+            print(org_forwarded.size())
+            print(scaled_lora_output.size())
+            raise e
        return x

    def enable_gradient_checkpointing(self: Module):