fix #95

* fixed mistake
wkpark · Dec 28, 2023 · 13e4b15 · 13e4b15
1 parent 73db45e
commit 13e4b15
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/scripts/rebasin/weight_matching.py b/scripts/rebasin/weight_matching.py
@@ -829,8 +829,8 @@ def scipylap(cost):
 
 
 def _valid_key(key):
-    if "cond_stage_model.transformer.text_model." in key:
-        return True
+    if "cond_stage_model." in key and "transformer.text_model." not in key:
+        return False
     return "model_" not in key