define_train_unet

n_channels = 3
n_classes = 2
# Next creat an instance of the UNet model
modelUnet = UNet().to(device)

criterion = torch.nn.BCEWithLogitsLoss()

# Now define the optimizer
optimizerUnet = optim.Adam(modelUnet.parameters(), lr = 0.00001, weight_decay=0)

# And finally lets train the model
# train(modelUnet, optimizerUnet, focal_loss, 500, dataloader2, print_status=True)
model_out = train(modelUnet, optimizerUnet, criterion, 200, dataloader2, print_status=True)