Might it be possible to import and export models in ONNX (https://onnx.ai/) format? In cases where one needs very high performance in the training phase, one can use an optimized library. Are developers of the libraries you are considering to integrate supportive?