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