diff --git a/instill/helpers/cli.py b/instill/helpers/cli.py index d41d3f9..b1a0483 100644 --- a/instill/helpers/cli.py +++ b/instill/helpers/cli.py @@ -118,6 +118,14 @@ def cli(): action="store_true", required=False, ) + run_parser.add_argument( + "-ng", + "--num-of-gpus", + help="number of gpus to use if gpu flag is on, default to 1", + type=int, + default=1, + required=False, + ) run_parser.add_argument( "-t", "--tag", @@ -311,7 +319,7 @@ def run(args): {args.name}:{args.tag} /bin/bash -c \ \"serve build _model:entrypoint -o serve.yaml && \ sed -i 's/app1/default/' serve.yaml && \ - sed -i 's/num_cpus: 0.0/num_gpus: 1.0/' serve.yaml && \ + sed -i 's/num_cpus: 0.0/num_gpus: {args.num_of_gpus}/' serve.yaml && \ serve run serve.yaml\"", shell=True, check=True,