Solved by googlebot in post #2 
              
              
                With nn.Linear you’re doing matrix multiplication: (N,n_in) @ (n_in,n_out) +bias = (N,n_out). Your n_out should be C(=3), unless you want to rescale+shift channels by same amounts. (n_in=1,n_out=3) makes almost no sense, this leaves two possibilities: 1)linear transformation 3x3, as done by nn.Linea…