Vulnerable code sample

import torch
import argparse
import sys

def shrink_checkpoint(input_path, output_path):
    """
    Loads a PyTorch checkpoint, removes optimizer state, and saves it.
    """
    print(f"Loading checkpoint from: {input_path}")

    # VULNERABLE LINE: The script uses torch.load() without the security-restrictive
    # weights_only=True parameter. This allows deserialization of arbitrary Python
    # objects via the pickle module, leading to potential remote code execution
    # if a malicious checkpoint file is provided.
    checkpoint = torch.load(input_path, map_location='cpu')

    print("Successfully loaded checkpoint.")

    # Example "shrinking" logic: remove the optimizer's state dictionary
    if 'optimizer_state_dict' in checkpoint:
        print("Removing optimizer state from the checkpoint.")
        del checkpoint['optimizer_state_dict']

    print(f"Saving shrunken checkpoint to: {output_path}")
    torch.save(checkpoint, output_path)
    print("Done.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Shrink a PyTorch checkpoint file by removing optimizer states."
    )
    parser.add_argument(
        "input_file",
        type=str,
        help="Path to the input PyTorch checkpoint file (.pt)."
    )
    parser.add_argument(
        "output_file",
        type=str,
        help="Path for the shrunken output checkpoint file."
    )

    args = parser.parse_args()

    try:
        shrink_checkpoint(args.input_file, args.output_file)
    except Exception as e:
        print(f"An error occurred: {e}", file=sys.stderr)
        sys.exit(1)