    def create_ait_module_wrapper(engine, interp_result, trace_ait_module, *inputs):
        """
        Some use cases need to torch.jit.script a model with AITModules in
        it, but TorchScript does not support variadic inputs. We can get
        around this by scripting the AITModule with some sample inputs.
        This is turned in by passing allow_scripting=True.
        """
        mod = AITModule(engine, interp_result)
        # sanity test before tracing
        mod(*inputs)
        return torch.jit.trace(mod, inputs) if trace_ait_module else mod
Some use cases need to torch.jit.script a model with AITModules in
it, but TorchScript does not support variadic inputs. We can get
around this by scripting the AITModule with some sample inputs.
This is turned in by passing allow_scripting=True.
fx2ait/fx2ait/ait_module.py