Okay ignore the second suggestion I made. I think it resulted in that error when using CUDA.
I think your loss function is actually working, but the gradients as a result are incredibly small. In my tests I’m seeing gradients as small as 1e-8. Which when using a learning rate such as 1e-2, the update to your weights are incredibly small. So your loss is appearing to be unmoving, but it actually is. Set your lambd to some larger number for testing, such as 10.0. I think you’ll see a change in your loss now.