import torch
import torchvision.models as models
import os

def download_resnet50_pretrained(save_dir):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    resnet50 = models.resnet50(pretrained=True)

    save_path = os.path.join(save_dir, 'resnet50_imagenet.pth')
    torch.save(resnet50.state_dict(), save_path)
    print(f"Model saved at {save_path}")

if __name__ == "__main__":
    save_directory = "./save/models/resnet50/"
    download_resnet50_pretrained(save_directory)
