I implemented a weighted NLLLoss2d in this thread. Maybe you could use it by setting your weights accordingly to 0 and 1.
