File tree Expand file tree Collapse file tree 1 file changed +9
-1
lines changed Expand file tree Collapse file tree 1 file changed +9
-1
lines changed Original file line number Diff line number Diff line change @@ -118,6 +118,14 @@ def cli():
118
118
action = "store_true" ,
119
119
required = False ,
120
120
)
121
+ run_parser .add_argument (
122
+ "-ng" ,
123
+ "--num-of-gpus" ,
124
+ help = "number of gpus to use if gpu flag is on, default to 1" ,
125
+ type = int ,
126
+ default = 1 ,
127
+ required = False ,
128
+ )
121
129
run_parser .add_argument (
122
130
"-t" ,
123
131
"--tag" ,
@@ -311,7 +319,7 @@ def run(args):
311
319
{ args .name } :{ args .tag } /bin/bash -c \
312
320
\" serve build _model:entrypoint -o serve.yaml && \
313
321
sed -i 's/app1/default/' serve.yaml && \
314
- sed -i 's/num_cpus: 0.0/num_gpus: 1.0 /' serve.yaml && \
322
+ sed -i 's/num_cpus: 0.0/num_gpus: { args . num_of_gpus } /' serve.yaml && \
315
323
serve run serve.yaml\" " ,
316
324
shell = True ,
317
325
check = True ,
You can’t perform that action at this time.
0 commit comments