If one defines a model f, and x is training data, then we have f(x).
How can one concisely define a model f(w) as a function of its parameters w? Will PyTorch be able to differentiate it wrt w, like JAX/STAX can?