with torch.inference_mode():
    traced = torch.jit.trace(
        model,
        trace_input,
        check_inputs=check_inputs,
    )

torch.jit.save(traced, artifact)