Problems with loading pre-trained model

Hey @hanikevi ,

I got a bug when I want to load a pretrained model from Lisas first validations. I think this might be related to our changes in opt tools to better structure the imports. Here is the bug:

jaxlib.xla_extension.XlaRuntimeError: Unknown custom type in pickled PyTreeDef: <class 'opt_tools.jax_tools.sysvars.TrajBatchFrankaData'>

Reproducable in main by running python -m tests.validation_real_data.

What do you mean on this?

Thanks!