search

shouxieai / infer

A new tensorrt integrate. Easy to integrate many tasks
MIT License
388 stars 80 forks source link

How to select which GPU to use? #11

Closed ghmchen closed 1 year ago

ghmchen commented 1 year ago

I have two GPUs. It would be great if I can point to which GPU to use by "auto yolo = yolo::load(modelpath, modelType, gpuID); " I am trying to modify codes by add "SetDevice(const int gpuID)" to the function construct() at infer.cn. But without sucess. Any sugestion? Thanks.

shouxieai commented 1 year ago

You can use the cudaSetDevice function before loading the model.

ghmchen commented 1 year ago

I will try it. Thanks.