Allow specifying device to use when training

Closes #139 (closed)

dan/datasets/extract/__init__.py

@@ -147,6 +147,13 @@ def add_extract_parser(subcommands) -> None:
         help="Images larger than this height will be resized to this height.",
     )
 
+    parser.add_argument(
+        "--subword-vocab-size",
+        type=int,
+        default=1000,
+        help="Size of the vocabulary to train the sentencepiece subword tokenizer needed for language model.",
+    )
+
     # Formatting arguments
     parser.add_argument(
         "--image-format",