How can use TRT in torchrec?

[yjjinjie]

I see the example: https://github.com/pytorch/torchrec/blob/v0.2.0/examples/inference/dlrm_predict_single_gpu.py

but how can i split the embedding and dense model,and then dense model use trt, the conat (ebc+ trt_dense model) , export the torchscript model in C++ inference ?

[PaulZhang12]

@yjjinjie that inference solution is legacy, will need to clean up. Check out this file for exporting a model to TorchScript for C++ serving: https://github.com/pytorch/torchrec/blob/main/torchrec/inference/dlrm_predict.py

You can apply lowering to the dense part of the model individually and then still torchscript it like the example above

[yjjinjie]

can you show a example--> how to lowering to the dense part of the model individually?

just lowering the denseArch?

https://github.com/pytorch/torchrec/blob/main/torchrec/models/dlrm.py#L115