After each epoch y_pred simply keeps increasing input at each batch is 64x10 tensor, trying to predict max of the vector at each row. I thought the gradient might not be going t
