exception for unknown arg should have been TypeError

This commit is contained in:
Ross Wightman 2024-08-12 16:09:36 -07:00
parent bd0f79153b
commit a7b0bfc376
3 changed files with 3 additions and 3 deletions

View File

@ -185,7 +185,7 @@ def load_pretrained(
check_hash=_CHECK_HASH, check_hash=_CHECK_HASH,
weights_only=True, weights_only=True,
) )
except ValueError: except TypeError:
state_dict = load_state_dict_from_url( state_dict = load_state_dict_from_url(
pretrained_loc, pretrained_loc,
map_location='cpu', map_location='cpu',

View File

@ -54,7 +54,7 @@ def load_state_dict(
else: else:
try: try:
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=weights_only) checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=weights_only)
except ValueError: except TypeError:
checkpoint = torch.load(checkpoint_path, map_location=device) checkpoint = torch.load(checkpoint_path, map_location=device)
state_dict_key = '' state_dict_key = ''

View File

@ -193,7 +193,7 @@ def load_state_dict_from_hf(
_logger.debug(f"[{model_id}] Safe alternative not found for '{filename}'. Loading weights using default pytorch.") _logger.debug(f"[{model_id}] Safe alternative not found for '{filename}'. Loading weights using default pytorch.")
try: try:
state_dict = torch.load(cached_file, map_location='cpu', weights_only=weights_only) state_dict = torch.load(cached_file, map_location='cpu', weights_only=weights_only)
except ValueError: except TypeError:
state_dict = torch.load(cached_file, map_location='cpu') state_dict = torch.load(cached_file, map_location='cpu')
return state_dict return state_dict