fix argparse bug

tsinghua-fib-lab · Dec 2, 2024 · 1f66170 · 1f66170
1 parent db41a69
commit 1f66170
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/main.cu b/src/main.cu
@@ -129,8 +129,9 @@ int main(int argc, char** argv) {
   }
 
   uint device = 0;
-  auto device_str = parser.get<std::string>("gpu");
-  device = std::stoi(device_str);
+  if (cfg["gpu"]) {
+    device = cfg["gpu"].as<int>();
+  }
 
   s.Init(name, {
                    .map_file = map_file,