changed the loss function a bit, and made the input variables a bit more distinct. Linear case works now