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