add no_cuda args in extract_features

27ee0fff · thomwolf · aa50fd19 · 27ee0fff
Commit 27ee0fff authored Nov 17, 2018 by thomwolf
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

examples/extract_features.py examples/extract_features.py +4 -0

No files found.
--- a/examples/extract_features.py
+++ b/examples/extract_features.py
@@ -208,6 +208,10 @@ def main():
                        type=int,
                        default=-1,
                        help = "local_rank for distributed training on gpus")
+    parser.add_argument("--no_cuda",
+                        default=False,
+                        action='store_true',
+                        help="Whether not to use CUDA when available")
    args = parser.parse_args()