Solved by ncuxomun in post #8 
              
              
                Hi. Yes, indeed, it did help. So, what I did was I had one Linear layer with a single output for x, which is then connected to a loss function (x_hat, x) and to another that has to make sure that (x_hat + y, ones(size(x_hat+y)). So, in total I had three loss functions. At last, I had to weigh each r…