From 4fe5bbe2cdaabca4a0d13da28f5a3649d33dcc6f Mon Sep 17 00:00:00 2001
From: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Date: Tue, 13 Aug 2024 14:07:39 -0700
Subject: [PATCH] Fix cache directory specification during testing (#786)

### Description

This PR fixes the cache directory specification when creating test
models.

### Motivation and Context

Previously, the provided cache directory would not be used when
`num_hidden_layers` was specified.
---
 src/python/py/models/builder.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 097e45be5..fb6523b33 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -1611,8 +1611,8 @@ def make_model(self, input_path):
             model = QuantModel.from_pretrained(self.quant_type, input_path, self.quant_attrs["bits"], self.quant_attrs["group_size"], self.quant_attrs["use_g_idx"], q_size, kv_size, self.intermediate_size, self.num_layers)
         else:
             # Load PyTorch model
-            extra_kwargs = {} if os.path.exists(self.model_name_or_path) else {"num_hidden_layers": self.num_layers} if "num_hidden_layers" in self.extra_options else {"cache_dir": self.cache_dir}
-            model = AutoModelForCausalLM.from_pretrained(self.model_name_or_path, use_auth_token=True, trust_remote_code=True, **extra_kwargs)
+            extra_kwargs = {"num_hidden_layers": self.num_layers} if "num_hidden_layers" in self.extra_options else {}
+            model = AutoModelForCausalLM.from_pretrained(self.model_name_or_path, cache_dir=self.cache_dir, use_auth_token=True, trust_remote_code=True, **extra_kwargs)
 
         # Loop through model and map each module to ONNX/ORT ops
         self.layer_id = 0