n1ck-guo commited on
Commit
048e7b0
1 Parent(s): 3785f47

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
@@ -147,8 +147,9 @@ def get_quantized_model_parameters_memory(model_info: ModelInfo, quant_method=""
147
  print(str(e))
148
 
149
  filenames = [sib.rfilename for sib in model_info.siblings]
150
- if "pytorch_model.bin" in filenames:
151
- url = hf_hub_url(model_info.id, filename="pytorch_model.bin")
 
152
  meta = get_hf_file_metadata(url)
153
  params_b = round(meta.size * 2 / 1e9, 2)
154
  size_gb = round(meta.size / 1e9, 2)
 
147
  print(str(e))
148
 
149
  filenames = [sib.rfilename for sib in model_info.siblings]
150
+ if "pytorch_model.bin" in filenames or "model.safetensors" in filenames:
151
+ bin_filename = "pytorch_model.bin" if "pytorch_model.bin" in filenames else "model.safetensors"
152
+ url = hf_hub_url(model_info.id, filename=bin_filename)
153
  meta = get_hf_file_metadata(url)
154
  params_b = round(meta.size * 2 / 1e9, 2)
155
  size_gb = round(meta.size / 1e9, 2)