Update src/transformers/models/vitpose/convert_vitpose_to_hf.py

Co-authored-by: Pavel Iakubovskii <[email protected]>
NielsRogge · Nov 20, 2024 · e588f4f · e588f4f
1 parent d8e6e2e
commit e588f4f
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/transformers/models/vitpose/convert_vitpose_to_hf.py b/src/transformers/models/vitpose/convert_vitpose_to_hf.py
@@ -306,6 +306,10 @@ def write_model(model_path, model_name, push_to_hub):
         raise ValueError("Model not supported")
     print("Conversion successfully done.")
 
+    # save the model to a local directory
+    model.save_pretrained(model_path)
+    image_processor.save_pretrained(model_path)
+
     if push_to_hub:
         print(f"Pushing model and image processor for {model_name} to hub")
         model.push_to_hub(f"danelcsb/{model_name}")