From fd0b2164b3865a61b09a4abdfead98ba458c7d44 Mon Sep 17 00:00:00 2001 From: Travis Addair Date: Tue, 9 Apr 2024 09:58:36 -0700 Subject: [PATCH] Fix import of EntryNotFoundError (#401) --- server/lorax_server/utils/weights.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/server/lorax_server/utils/weights.py b/server/lorax_server/utils/weights.py index 1bc777d92..1779aa074 100644 --- a/server/lorax_server/utils/weights.py +++ b/server/lorax_server/utils/weights.py @@ -7,7 +7,7 @@ import torch import torch.distributed from huggingface_hub import hf_hub_download -from huggingface_hub.utils import LocalEntryNotFoundError +from huggingface_hub.utils import EntryNotFoundError, LocalEntryNotFoundError from loguru import logger from safetensors import SafetensorError, safe_open @@ -422,7 +422,7 @@ def download_weights( model_source.download_model_assets() return # No weights found on the hub with this extension - except utils.EntryNotFoundError as e: + except EntryNotFoundError as e: # Check if we want to automatically convert to safetensors or if we can use .bin weights instead if not extension == ".safetensors" or not auto_convert: raise e