fix bug
Browse filesSigned-off-by: n1ck-guo <heng.guo@intel.com>
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 |
-
|
|
|
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)
|