Ignore None values in converter

2015aroras · Oct 29, 2024 · 75dabae · 75dabae
1 parent 30844ec
commit 75dabae
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/transformers/models/olmo/convert_olmo_weights_to_hf.py b/src/transformers/models/olmo/convert_olmo_weights_to_hf.py
@@ -136,6 +136,8 @@ def write_model(
         state_dict[f"model.layers.{layer_i}.self_attn.rotary_emb.inv_freq"] = inv_freq
 
         for k, v in state_dict.items():
+            if v is None:
+                continue
             index_dict["weight_map"][k] = filename
             param_count += v.numel()
         torch.save(state_dict, os.path.join(tmp_model_path, filename))
@@ -153,6 +155,8 @@ def write_model(
     }
 
     for k, v in state_dict.items():
+        if v is None:
+            continue
         index_dict["weight_map"][k] = filename
         param_count += v.numel()
     torch.save(state_dict, os.path.join(tmp_model_path, filename))