refine for cuda graph (#301)

microsoft · Apr 23, 2024 · 0f7fb7a · 0f7fb7a
1 parent b1180a6
commit 0f7fb7a
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/src/models/model.cpp b/src/models/model.cpp
@@ -544,12 +544,9 @@ void Model::GetMaxBatchSizeFromGeneratorParams(const GeneratorParams& params) {
   max_batch_size_ = params.max_batch_size;
 
   if (DeviceType::CUDA == device_type_) {
-    if (max_batch_size_ == 0 && is_cuda_graph_enabled) {
-      throw std::runtime_error("CUDA graph is enabled, but max_batch_size is not set.");
-    }
-    if (max_batch_size_ > 0) {
-      if (!is_cuda_graph_enabled) {
-        throw std::runtime_error("CUDA graph is not enabled.");
+    if (is_cuda_graph_enabled) {
+      if (max_batch_size_ == 0) {
+        throw std::runtime_error("CUDA graph is enabled, but max_batch_size is not set.");
       }
       use_cuda_graph_ = true;
     }