Geometry optimisation with DM21 in PySCF?

[Kirill Kulaev]

Is it available to use DM21 for geometry optimization of molecule using PySCF?
It works only with HF nuclear gradients like as:

mf = scf.RKS(mol)
mf._numint = dm21.NeuralNumInt(dm21.Functional.DM21)
mf.apply(grad.rhf.Gradients).as_scanner()

But doesn't work with RKS gradients in pyscf.grad.rks.Gradients
AttributeError: 'NoneType' object has no attribute 'mol' in neural_numint.py in eval_xc()
How can I use RKS nuclear gradients?