import torch.optim as optim
optimizer = optim.SGD(unet.parameters(),
                      lr=0.01,
                      momentum=0.9,
                      weight_decay=0.0005)
    
criterion = nn.BCELoss()

Here is my optimizer