def fine_tune_model(num_epochs, data_dir, learning_rate, momentum):
    'Load a pretrained model and reset the final fully connected layer.'
    run.log('lr', np.float(learning_rate))
    run.log('momentum', np.float(momentum))
    model_ft = models.resnet18(pretrained=True)
    num_ftrs = model_ft.fc.in_features
    model_ft.fc = nn.Linear(num_ftrs, 2)
    device = torch.device(('cuda:0' if torch.cuda.is_available() else 'cpu'))
    model_ft = model_ft.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer_ft = optim.SGD(model_ft.parameters(), lr=learning_rate, momentum=momentum)
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
    model = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs, data_dir)
    return model
Load a pretrained model and reset the final fully connected layer.
def fine_tune_model(num_epochs, data_dir, learning_rate, momentum):
    
    run.log('lr', np.float(learning_rate))
    run.log('momentum', np.float(momentum))
    model_ft = models.resnet18(pretrained=True)
    num_ftrs = model_ft.fc.in_features
    model_ft.fc = nn.Linear(num_ftrs, 2)
    device = torch.device(('cuda:0' if torch.cuda.is_available() else 'cpu'))
    model_ft = model_ft.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer_ft = optim.SGD(model_ft.parameters(), lr=learning_rate, momentum=momentum)
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
    model = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs, data_dir)
    return model