n1ck-guo commited on
Commit
3785f47
1 Parent(s): 61c1bf5

Signed-off-by: n1ck-guo <heng.guo@intel.com>

Files changed (1) hide show
  1. src/submission/check_validity.py +3 -2
src/submission/check_validity.py CHANGED
@@ -154,8 +154,9 @@ def get_quantized_model_parameters_memory(model_info: ModelInfo, quant_method=""
154
  size_gb = round(meta.size / 1e9, 2)
155
  return params_b, size_gb
156
 
157
- if "pytorch_model.bin.index.json" in filenames:
158
- index_path = hf_hub_download(model_info.id, filename="pytorch_model.bin.index.json")
 
159
  """
160
  {
161
  "metadata": {
 
154
  size_gb = round(meta.size / 1e9, 2)
155
  return params_b, size_gb
156
 
157
+ if "pytorch_model.bin.index.json" in filenames or "model.safetensors.index.json" in filenames:
158
+ json_file_name = "pytorch_model.bin.index.json" if "pytorch_model.bin.index.json" in filenames else "model.safetensors.index.json"
159
+ index_path = hf_hub_download(model_info.id, filename=json_file_name)
160
  """
161
  {
162
  "metadata": {