Merge pull request #110 from BlackSamorez/pretrained_fix

find_predefined_tensor_parallel_config try-except fix
BlackSamorez · Jul 26, 2023 · c40153c · c40153c
2 parents 31c7097 + f000c03
commit c40153c
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/setup.cfg b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = tensor_parallel
-version = 1.3.0
+version = 1.3.1
 author = Andrei Panferov and Yaroslav Lisnyak
 author_email = [email protected]
 description = Automatically shard your large model between multiple GPUs, works without torch.distributed

diff --git a/src/tensor_parallel/pretrained_model.py b/src/tensor_parallel/pretrained_model.py
@@ -24,9 +24,9 @@ def find_predefined_tensor_parallel_config(
 ) -> Optional[Config]:
     device_ids = check_device_ids(device_ids)
 
-    try:
+    if model_config.model_type in PREDEFINED_CONFIGS:
         return PREDEFINED_CONFIGS[model_config.model_type](model_config, device_ids)
-    except KeyError:
+    else:
         logger.warning(
             "Using automatic config: tensor parallel config not provided and no custom config registered for the model"
         )