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
-