Just a quick idea. You could define a custom model Module that unpacks x, y and w, and sends x through your original model. For the output, combine w with model(x). The custom loss function then unpacks w from x and returns the loss.
Just an idea - the proof is in the doing, which I have not tried.