Improve get_mapped_key heuristic (#1525)

pytorch · Sep 9, 2024 · 66590b4 · 66590b4
1 parent 68d4f3e
commit 66590b4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torchtune/models/convert_weights.py b/torchtune/models/convert_weights.py
@@ -47,7 +47,8 @@
 
 def get_mapped_key(key: str, mapping_dict: Dict[str, str]) -> str:
     try:
-        if "layers" in key:
+        # Checks if there is a layer # in the key
+        if any(k.isdigit() for k in key.split(".")):
             # Replace layer number with "{}" to create key for lookup
             abstract_key = re.sub(r"(\.\d+)", ".{}", key)
             layer_num = re.search(r"\d+", key).group(0)