Function torch_tensorrt::set_device ¶
-
Defined in File torch_tensorrt.h
Function Documentation ¶
-
TORCHTRT_API
void
torch_tensorrt
::
set_device
(
const
int
gpu_id
)
¶
-
Set gpu device id.
- Parameters
-
gpu_id – Sets gpu id using cudaSetDevice