[Task] Research/prototype torch module for Torchscript compatibility

[edknv]

Description

The team had discussions on porting Transformers4rec to Merlin Models and supporting the model in Triton inference server as one of the goals for 22.09. However, there needs to be some de-risking before we are able to fully commit, because the pytorch module of Merlin Models is much less developed than the tensorflow counterpart and we have encountered some issues in the past with running Torchscript for Transformers4rec.

The goal of this task is to fully understand the risk around Torchscript and the effort involved in supporting session-based models in pytorch in Models. The task is to

• Start a prototype of the torch code by porting over the existing Tensorflow code to torch and write unit tests that make sure we can run torch.jit.script and/or torch.jit.trace,
• And also from the other side, run the existing Transformers4rec through Torchscript and understand the obstacles.

Starting point is translating tensorflow code for the main blocks SequentialBlock, ParallelBlock, and MLPBlock, and writing unit tests for torch.jit for these blocks. @marcromeyn has already started some work in porting the TF blocks to torch, which is available in the torch/clean-up branch.

Additional context

• Marc's branch: https://github.com/NVIDIA-Merlin/models/compare/main...torch/clean-up
• Past issues: https://github.com/NVIDIA-Merlin/Transformers4Rec/issues/282
• https://github.com/triton-inference-server/pytorch_backend