loss fails in backward

[skanderbug]

Hello, I find ImageGradientLoss in loss/loss.py may not be able to perform backward since
pred = pred.argmax(1).view(size[0], 1, size[2], size[3]).float()
is not differentiable.

I checked it in some toy experiment and maybe we can use
pred = pred [:, 1, :, :]
instead.

[wonbeomjang]

Thank you for letting me know. I will do an experiment soon, and I will reflect on the actual performance when it.