How to get the parameter of DualPotential solved by W2NeuralDual?

[genzm]

Hello! Thank you for the amazing library!!
And sorry for the very basic question.
I want to ask how to get the parameter of DualPotential which is solved by W2NeuralDual.
As written in document, I understand W2NeuralDual is computing optimal parameter $\theta$ for dual OT Kantorovich problem. (https://ott-jax.readthedocs.io/en/latest/solvers/_autosummary/ott.solvers.nn.neuraldual.W2NeuralDual.html#ott.solvers.nn.neuraldual.W2NeuralDual)
My question is, is there some way to get the optimal parameter $\theta$ from the learned potential?