From ab1a65f13a1935d089d38bff514c22722bb9225c Mon Sep 17 00:00:00 2001
From: Gnome Ann <>
Date: Mon, 15 Nov 2021 11:56:02 -0500
Subject: [PATCH] Fix hidden size detection for GPTJForCausalLM

---
 aiserver.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/aiserver.py b/aiserver.py
index 18fe309b..b4e8716b 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -600,17 +600,20 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
             try:
                 return int(model.transformer.hidden_size)
             except:
-                return int(model.transformer.embed_dim)
+                try:
+                    return int(model.transformer.embed_dim)
+                except:
+                    return int(model.lm_head.in_features)
 
         # If custom GPT Neo model was chosen
         if(vars.model == "NeoCustom"):
             model_config = open(vars.custmodpth + "/config.json", "r")
             js   = json.load(model_config)
-            vars.modeldim = int(js['hidden_size'])
             if("model_type" in js):
                 model     = AutoModelForCausalLM.from_pretrained(vars.custmodpth)
             else:
                 model     = GPTNeoForCausalLM.from_pretrained(vars.custmodpth)
+            vars.modeldim = get_hidden_size_from_model(model)
             tokenizer = GPT2Tokenizer.from_pretrained(vars.custmodpth)
             # Is CUDA available? If so, use GPU, otherwise fall back to CPU
             if(vars.hascuda):