fix: forgot to always set _disable_torch_cuda_device_set

Signed-off-by: Terry Kong <terryk@nvidia.com>

nemo/export/trt_llm/tensorrt_llm_run.py

@@ -502,6 +502,7 @@ def load_distributed(engine_dir, model_parallel_rank, gpus_per_node):
  # We want the engine to have the mp_rank, but the python runtime to not resassign the device of the current process
  # So we will set it to the current
  rank=torch.cuda.current_device(),
+ _disable_torch_cuda_device_set=True,
  )
 
  tensorrt_llm_worker_context.decoder = decoder