diff --git a/benchmarks/torchbench_model.py b/benchmarks/torchbench_model.py index fd30ee359871..6fa64b0329ac 100644 --- a/benchmarks/torchbench_model.py +++ b/benchmarks/torchbench_model.py @@ -76,6 +76,7 @@ { "test": "eval", "xla": "PJRT", + "dynamo": None, }, # TIMEOUT ], "hf_T5_generate": [ @@ -85,6 +86,7 @@ { "test": "eval", "xla": "PJRT", + "dynamo": None, }, # TIMEOUT ], "doctr_det_predictor": [{ @@ -135,6 +137,13 @@ ], } +# Models that had more graphs to be compiled than the actual size of +# the cache. +NEED_LARGER_CACHE = { + "cm3leon_generate", + "hf_T5_generate", +} + class TorchBenchModelLoader(ModelLoader): @@ -401,6 +410,9 @@ def update_process_env(self, process_env): if precision_flag is not None: process_env[precision_flag] = '1' + if self.model_name in NEED_LARGER_CACHE: + process_env["XLA_COMPILATION_CACHE_SIZE"] = "2048" + def pick_grad(self): # special case if self.model_name in ("maml",):