Update finetune_gradio.py

The safest approach here is to explicitly set weights_only=True to load only the model weights and avoid executing potentially unsafe code
This commit is contained in:
mahwiz khalil
2025-01-24 00:31:53 -08:00
committed by GitHub
parent 129014c5b4
commit c54f4e7fc0

View File

@@ -974,7 +974,7 @@ def calculate_train(
def extract_and_save_ema_model(checkpoint_path: str, new_checkpoint_path: str, safetensors: bool) -> str:
try:
checkpoint = torch.load(checkpoint_path)
checkpoint = torch.load(checkpoint_path, weights_only=True)
print("Original Checkpoint Keys:", checkpoint.keys())
ema_model_state_dict = checkpoint.get("ema_model_state_dict", None)