Hi everyone.. I have a class like this:

class NERRunnable(bentoml.Runnable):
    SUPPORTED_RESOURCES = ("<http://nvidia.com/gpu|nvidia.com/gpu>",)
    SUPPORTS_CPU_MULTI_THREADING = True

    def __init__(self):
        <http://bentoml_logger.info|bentoml_logger.info>("NERRunnable: Starting runner...")
        self.is_ready = True
        available_gpus = os.getenv("CUDA_VISIBLE_DEVICES", "")
        if available_gpus not in ("", "-1"):
            # assign GPU resources
            if not available_gpus.isdigit():
                raise ValueError(
                    "Expecting numeric value for CUDA_VISIBLE_DEVICES, got"
                    f" {available_gpus}."
                )
        else:
            raise ValueError("No GPU resources available.")
        self.device = int(available_gpus)
        tokenizer = AutoTokenizer.from_pretrained(settings.MODEL_NAME_FOLDER)
        model = AutoModelForTokenClassification.from_pretrained(
            settings.MODEL_NAME_FOLDER
        )
        self.pipe = pipeline(
            task="ner", model=model, tokenizer=tokenizer, device=self.device
        )
        <http://bentoml_logger.info|bentoml_logger.info>("NERRunnable: Compiling model...")
        self.pipe.model = torch.compile(model, mode="reduce-overhead")
        self.pre_heat_result = self.pipe(
            ["pre heat model"], aggregation_strategy="average"
        )
        <http://bentoml_logger.info|bentoml_logger.info>("NERRunnable: Runner is READY...")
        <http://bentoml_logger.info|bentoml_logger.info>("NERRunnable: %s", len(self.pre_heat_result))
        self.is_ready = True

    @bentoml.Runnable.method(batchable=True, batch_dim=0)
    def get_entities(self, sentences: List[str]):
        # Tokenize sentences
        return self.pipe(sentences, aggregation_strategy="average")

When I remove the:

self.pre_heat_result = self.pipe(
            ["pre heat model"], aggregation_strategy="average"
        )

Everything works fine.. Otherwise I get something like:

File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/_aot_autograd/utils.py", line 81, in g
  return f(*args)
 File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 94, in runtime_wrapper
  all_outs = call_func_at_runtime_with_args(
 File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/_aot_autograd/utils.py", line 105, in call_func_at_runtime_with_args
  out = normalize_as_list(f(args))
 File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 118, in rng_functionalization_wrapper
  return compiled_fw(args)
 File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/codecache.py", line 864, in __call__
  return self.get_current_callable()(inputs)
 File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/compile_fx.py", line 665, in run
  return compiled_fn(new_inputs)
 File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/cudagraph_trees.py", line 380, in deferred_cudagraphify
  fn, out = cudagraphify(model, inputs, new_static_input_idxs, *args, **kwargs)
 File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/cudagraph_trees.py", line 400, in cudagraphify
  manager = get_container(device_index).get_tree_manager()
 File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/cudagraph_trees.py", line 339, in get_container
  container_dict = get_obj(local, "tree_manager_containers")
 File "/usr/local/lib/python3.10/dist-packages/torch/_inductor/cudagraph_trees.py", line 334, in get_obj
  assert torch._C._is_key_in_tls(attr_name)
AssertionError

On my load tests, I notices that the first inference is always slower than expected.. Do you know how can skip that and use this torch.compile on bentoml ?