feat(main): Add use_cuda and use_dml parameters of CLI

RapidAI · Jun 25, 2024 · 0d88463 · 0d88463
1 parent 98c10ec
commit 0d88463
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/rapid_layout/main.py b/rapid_layout/main.py
@@ -158,6 +158,12 @@ def main():
         choices=list(KEY_TO_MODEL_URL.keys()),
         help="IoU threshold, the range is [0, 1]",
     )
+    parser.add_argument("--use_cuda", action="store_true", help="Whether to use cuda.")
+    parser.add_argument(
+        "--use_dml",
+        action="store_true",
+        help="Whether to use DirectML, which only works in Windows10+.",
+    )
     parser.add_argument(
         "-v",
         "--vis",