3D Unet Implementation doesn't overfit

vision
6

I’m able to overfit the model on random data using:

device = 'cuda'
model = UNet().to(device)
x = torch.randn(1, 1, 116, 132, 132).to(device)
target = torch.randint(0, 2, (1, 1, 28, 44, 44), device=device).float()

criterion = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

for idx in range(500):
    optimizer.zero_grad()
    output = model(x)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    print('Iter {}, loss {:.4f}'.format(idx, loss.item()))

In your code it seems you are somehow normalizing the labels.
Does it mean you don’t use a target tensor containing only zeros and ones?
If so, note that binary cross entropy will only yield a zero loss, if the target and predictions are either both (close to) zero or one as described here.