import os
import torch

def main():
    # make sure the directory exists
    checkpoint_dir = 'cifar100_models/state_dicts'
    os.makedirs(checkpoint_dir, exist_ok=True)

    model_list = [m for m in torch.hub.list("chenyaofo/pytorch-cifar-models", force_reload=True) if m.startswith("cifar100")]

    for mn in model_list:
        print(f"Loading model '{mn}'...")
        model = torch.hub.load("chenyaofo/pytorch-cifar-models", mn, pretrained=True)
        print(f"Model '{mn}' loaded successfully.")
        # save the model checkpoint
        checkpoint_path = os.path.join(checkpoint_dir, f'{mn}.pth')
        torch.save({'model': model.state_dict()}, checkpoint_path, _use_new_zipfile_serialization=False)
        print(f"Model checkpoint saved to '{checkpoint_path}'.")

def print_cifar100_models():
    print([m for m in torch.hub.list("chenyaofo/pytorch-cifar-models", force_reload=True) if m.startswith("cifar100")])

if __name__ == '__main__':
    main()
    print_cifar100_models()
    