[Fix] Fix embedding shape check in ChatModule (#1953)

This PR is a fix to address #1952.
mlc-ai · Mar 13, 2024 · 8d192ef · 8d192ef
1 parent 5c29f02
commit 8d192ef
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cpp/llm_chat.cc b/cpp/llm_chat.cc
@@ -1447,7 +1447,7 @@ class LLMChat {
       embedding_shape = embedding_nd.Shape();
     }
     ICHECK_EQ(embedding_shape.size(), 2);
-    ICHECK_GT(embedding_shape[0], 1);
+    ICHECK_GE(embedding_shape[0], 1);
     this->hidden_size_ = embedding_shape[1];
     return this->hidden_size_;
   }