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