From c692987e40cdc128c0138ab04f7f7b2ca6549409 Mon Sep 17 00:00:00 2001
From: Gnome Ann <>
Date: Fri, 20 May 2022 14:54:49 -0400
Subject: [PATCH] Fix an error that occurs when loading GPT-2 models I forgot
 that this new_rebuild_tensor function's first argument's type is different
 when loading GPT-2 models.

---
 aiserver.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/aiserver.py b/aiserver.py
index b9b9267b..b1fce2cb 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -1613,10 +1613,13 @@ if(not vars.use_colab_tpu and vars.model not in ["InferKit", "Colab", "OAI", "Go
                         model     = GPTNeoForCausalLM.from_pretrained("models/{}".format(vars.model.replace('/', '_')), revision=vars.revision, cache_dir="cache", **lowmem)
                 else:
                     old_rebuild_tensor = torch._utils._rebuild_tensor
-                    def new_rebuild_tensor(storage, storage_offset, shape, stride):
-                        dtype = storage.storage_type.dtype
-                        if(not isinstance(dtype, torch.dtype)):
-                            dtype = storage.storage_type(0).dtype
+                    def new_rebuild_tensor(storage: Union[torch_lazy_loader.LazyTensor, torch.Storage], storage_offset, shape, stride):
+                        if(not isinstance(storage, torch_lazy_loader.LazyTensor)):
+                            dtype = storage.dtype
+                        else:
+                            dtype = storage.storage_type.dtype
+                            if(not isinstance(dtype, torch.dtype)):
+                                dtype = storage.storage_type(0).dtype
                         if(dtype is torch.float32 and len(shape) >= 2):
                             vars.fp32_model = True
                         return old_rebuild_tensor(storage, storage_offset, shape, stride)