fix: ignore tensors with the particular dim while loading (leejet#233)

phyllispeng123 · Apr 29, 2024 · 760cfaa · 760cfaa
1 parent 6d16f68
commit 760cfaa
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/model.cpp b/model.cpp
@@ -888,6 +888,11 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
             }
         }
 
+        // ggml/src/ggml.c:2745
+        if (n_dims < 1 || n_dims > GGML_MAX_DIMS) {
+            continue;
+        }
+
         TensorStorage tensor_storage(prefix + name, type, ne, n_dims, file_index, ST_HEADER_SIZE_LEN + header_size_ + begin);
         tensor_storage.reverse_ne();