Function torch_tensorrt::set_device

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