Custom loss function no_grad and OOM error

I want to use a custom loss function, where I do a compute condition before using the MSELoss what ever I do getting error no_grad if I used the requires_grad=True I get OOM error, note my batch size was working fine before the custom loss_fn, when I reduce the batch size to 2 it works but the loss doesn’t improve at all, for about 10k steps nothing changes same value.

pred = model(input)
loss_fn = nn.MSELoss()
diff =, target), 0).type(pred.type())
tar = torch.ones(diff.size()).type(diff.type()).cuda()
loss = loss_fn(diff, tar)

I think “diff” is True or False and “tar” is float. MSE might give error. Isn’t it? Can you post “tar” and “loss” value if its not giving error?

I have changed the diff type by .type(pred.type()) to float, I don’t think that’s the problem, it’s like the training graph is detached when using the and torch.sub is not differentiable.

I also try using math operations ((pred-target)>0).float() but still the same error

Greater than itself is non differentiable. You should not use it if you want gradient flow.

so there’s no solution for me to use this loss function, without this step my loss function it’s not complete