modified https://gist.github.com/joecummings/60f413cc0400a501d55c8672d8b5b393 to work with tracing

kwargs.py

import torch
import torch.nn as nn

# LM example
class EncoderDecoderLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()

    @torch.jit.export
    def prepare_inputs(self, var1: torch.Tensor, var2: torch.Tensor):
        return {
          "var1": var1,
          "var2": var2,
        }

    def forward(self, var1 = None, var2 = None):
        return var1*var2

# Parent class example
class Parent(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = torch.jit.script(model)
        self.x = torch.rand((2, 2))
        self.y = torch.rand((2, 2))

    def forward(self):
        kwargs = self.model.prepare_inputs(self.x, self.y)
        return self.model(**kwargs) ### <---- this is the issue

# Initialization of Torchscriptable model
lm = EncoderDecoderLanguageModel()

# # Initialization of parent class that will call forward on torchscriptable model
parent_class = Parent(model=lm)
parent_class_traced = torch.jit.trace(parent_class, ())