Update lora.py

hako-mikan · Jan 21, 2024 · bb7e90f · bb7e90f
1 parent e0a384d
commit bb7e90f
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/trainer/lora.py b/trainer/lora.py
@@ -344,11 +344,11 @@ def save_weights(self, file, t, metaname):
             torch.save(state_dict, file)
 
     def __enter__(self):
-        for lora in self.unet_loras:
+        for lora in self.unet_loras + self.te_loras:
             lora.multiplier = 1.0
 
     def __exit__(self, exc_type, exc_value, tb):
-        for lora in self.unet_loras:
+        for lora in self.unet_loras + self.te_loras:
             lora.multiplier = 0
 
     def check_weight(self):
@@ -708,11 +708,11 @@ def enumerate_params(loras):
         return all_params
 
     def __enter__(self):
-        for lora in self.unet_loras:
+        for lora in self.unet_loras + self.te_loras:
             lora.multiplier = 1.0
 
     def __exit__(self, exc_type, exc_value, tb):
-        for lora in self.unet_loras:
+        for lora in self.unet_loras + self.te_loras:
             lora.multiplier = 0
 
     def prepare_grad_etc(self):