Skip to content

Commit

Permalink
Fix cache directory specification during testing (#786)
Browse files Browse the repository at this point in the history
### 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.
  • Loading branch information
kunal-vaishnavi authored Aug 13, 2024
1 parent 99275a3 commit 4fe5bbe
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions src/python/py/models/builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 4fe5bbe

Please sign in to comment.