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
-