So, I have a lot of things going on under the hood, but everything is pretty normal, no fancy stuff going on. It should be a pretty normal looking codebase.
Note that you are not using any non-linearity between the conv layers.
Given that and a kernel of a single pixel, it looks like your model captures the blue-ish color successfully.
Yep! You are right! I added ReLU to get a better convergence.
But, still didn’t reached 0 loss. I increased the kernel_size upto 5x5 for both the Conv Layers. I went very close, but still didn’t converged. Shouldn’t the loss go to 0?
If the model has enough “capacity”, then the loss should converge towards zero.
Just for the sake of debugging this example, you could increase the number of kernels, their spatial size, or try a huge linear layer instead.