Tracing and saving Ensemble Custom C++ and CUDA extensions using torch.jit.save

Viewed 77

Summary: I'm trying to export a single model that uses multiple models built from CPP and other TorchScript files using python.

I've read LOADING A TORCHSCRIPT MODEL IN C++ and learnt that it is possible to save and load models in torchscript.

I've also read CUSTOM C++ AND CUDA EXTENSIONS and learnt I can use custom CPP compiled functions / extensions for optimizing my code.

During forward pass, my model is using a model1 (loaded from .pt file) for feature extraction, a function func from a CPP file for some post processing of the feature maps and model2 for decoding the post processed features.

Would the result of torch.jit.trace contain the "logic" of all the models and CPP functions used in my code?

0 Answers
Related