bug in env var USE_SMALL_MODELS (#236)

pull/247/head
Jairo Correa 2023-05-03 17:21:35 -03:00 committed by GitHub
parent 29e68a2e1b
commit dd17fb00af
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions

View File

@ -137,7 +137,7 @@ def _grab_best_device(use_gpu=True):
def _get_ckpt_path(model_type, use_small=False):
key = model_type
if use_small:
if use_small or USE_SMALL_MODELS:
key += "_small"
return os.path.join(CACHE_DIR, REMOTE_MODEL_PATHS[key]["file_name"])