I’m still unsure, if your model is returning logits or log probabilities or something else.
That being said, I think this line of code is not working as expected:
# Setup using your provided shapes
output = torch.randn(1, 2, 64, 64)
target = torch.randint(0, 2, (1, 64, 64))

# Overlap calculation
tmp = torch.round(output).mul(target)
print(tmp.shape)
> torch.Size([1, 2, 64, 64])

I guess you would tile to get the predicted class from output, which would work using torch.argmax(output, 1) instead of torch.round.
However, I might be wrong depending on the actual type of output you are dealing with.