Due to the large dataset, I have spend lots of time testing. How to speed up this process?
Main code:
# Some other function such as get args

def validate(valloader, model):
    scores = []
    with torch.no_grad():
        for i_iter, batch in enumerate(valloader):
            images, labels = batch
            images = images.cuda(non_blocking=True)

            preds = model(images)
            outputs = torch.nn.functional.softmax(preds)[:, -1]
            scores.extend(outputs.cpu().numpy())
    return scores

if __name__ == '__main__':
    args = parse_args()
    print("Calling with args:")
    print(args)

    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
    cudnn.enabled = True
    cudnn.benchmark = True
    torch.backends.cudnn.deterministic = False
    torch.backends.cudnn.enabled = True

    # Define dataset
    rdataset = MyDataSet(dataset='test', size=224, mode='both')
    pdataset = MyDataSet(dataset='val', size=224, mode='both')
    rloader = data.DataLoader(rdataset, batch_size=args.batch_size, shuffle=False,
                                  num_workers=1)  # , pin_memory=True)
    ploader = data.DataLoader(pdataset, batch_size=args.batch_size, shuffle=False,
                                num_workers=1)  # , pin_memory=True)

    # Define model
    num_classes = 2
    model = resnet50(pretrained=False, num_classes=num_classes)
    model.eval()
    model.cuda()

    prefix = args.model.split('/')
    pattern = prefix[-1]
    prefix = '/'.join(prefix[:-1])

    all_files = os.listdir(prefix)
    all_models = [os.path.join(prefix, f) for f in all_files if pattern in f]

    new_state_dict = model.state_dict().copy()

    for i in range(len(all_models)):
        try:
            old_state_dict = torch.load(all_models[i])
            new_state_dict = dict()
            for k, v in old_state_dict.items():
                new_state_dict[k[7:]] = v
            model.load_state_dict(new_state_dict)
        except:
            print("Fail when loading model %s" % all_models[i])
            continue

        rscores = validate(rloader, model)
        pscores = validate(ploader, model)

        # Other subprocess