From 9d5025f531a34c7ba0d44c342bff2b60834e6bc5 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Fri, 19 May 2023 11:20:08 -0300
Subject: [PATCH] Improve error handling while loading GPTQ models

---
 modules/GPTQ_loader.py | 2 +-
 modules/models.py      | 5 ++++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/modules/GPTQ_loader.py b/modules/GPTQ_loader.py
index 55c84ad..8fef5a8 100644
--- a/modules/GPTQ_loader.py
+++ b/modules/GPTQ_loader.py
@@ -140,7 +140,7 @@ def load_quantized(model_name):
     if shared.args.model_type is None:
         logging.error("The model could not be loaded because its type could not be inferred from its name.")
         logging.error("Please specify the type manually using the --model_type argument.")
-        return
+        return None
 
     # Select the appropriate load_quant function
     model_type = shared.args.model_type.lower()
diff --git a/modules/models.py b/modules/models.py
index a04a7ec..867f463 100644
--- a/modules/models.py
+++ b/modules/models.py
@@ -97,7 +97,10 @@ def load_model(model_name):
         model, tokenizer = output
     else:
         model = output
-        tokenizer = load_tokenizer(model_name, model)
+        if model is None:
+            return None, None
+        else:
+            tokenizer = load_tokenizer(model_name, model)
 
     # Hijack attention with xformers
     if any((shared.args.xformers, shared.args.sdp_attention)):