Setting specific device_id with set_current_gpu_device_id not working

[MadMenHitBooker]

Describe the bug
Hello. I tried to use a method to set the device_id, even if not documented. Sadly it works only on device 0

To Reproduce
This works

import onnxruntime_genai as og

# set device id
og.set_current_gpu_device_id(0)

This doesn't

import onnxruntime_genai as og

# set device id
og.set_current_gpu_device_id(1)

I get this error

2024-07-29 15:55:02.094196786 [E:onnxruntime:onnxruntime-genai, inference_session.cc:2045 operator()] Exception during initialization: /onnxruntime_src/onnxruntime/core/providers/cuda/cuda_call.cc:123 std::conditional_t<THRW, void, onnxruntime::common::Status> onnxruntime::CudaCall(ERRTYPE, const char*, const char*, ERRTYPE, const char*, const char*, int) [with ERRTYPE = cudnnStatus_t; bool THRW = true; std::conditional_t<THRW, void, onnxruntime::common::Status> = void] /onnxruntime_src/onnxruntime/core/providers/cuda/cuda_call.cc:116 std::conditional_t<THRW, void, onnxruntime::common::Status> onnxruntime::CudaCall(ERRTYPE, const char*, const char*, ERRTYPE, const char*, const char*, int) [with ERRTYPE = cudnnStatus_t; bool THRW = true; std::conditional_t<THRW, void, onnxruntime::common::Status> = void] CUDNN failure 7: CUDNN_STATUS_MAPPING_ERROR ; GPU=0 ; hostname=041ac24945f2 ; file=/onnxruntime_src/onnxruntime/core/providers/cuda/cuda_execution_provider.cc ; line=182 ; expr=cudnnSetStream(cudnn_handle_, stream); 


Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
onnxruntime_genai.onnxruntime_genai.OrtException: Exception during initialization: /onnxruntime_src/onnxruntime/core/providers/cuda/cuda_call.cc:123 std::conditional_t<THRW, void, onnxruntime::common::Status> onnxruntime::CudaCall(ERRTYPE, const char*, const char*, ERRTYPE, const char*, const char*, int) [with ERRTYPE = cudnnStatus_t; bool THRW = true; std::conditional_t<THRW, void, onnxruntime::common::Status> = void] /onnxruntime_src/onnxruntime/core/providers/cuda/cuda_call.cc:116 std::conditional_t<THRW, void, onnxruntime::common::Status> onnxruntime::CudaCall(ERRTYPE, const char*, const char*, ERRTYPE, const char*, const char*, int) [with ERRTYPE = cudnnStatus_t; bool THRW = true; std::conditional_t<THRW, void, onnxruntime::common::Status> = void] CUDNN failure 7: CUDNN_STATUS_MAPPING_ERROR ; GPU=0 ; hostname=041ac24945f2 ; file=/onnxruntime_src/onnxruntime/core/providers/cuda/cuda_execution_provider.cc ; line=182 ; expr=cudnnSetStream(cudnn_handle_, stream); 

[baijumeswani]

@MadMenHitBooker what onnxruntime-genai package are you using? I was not able to reproduce the error reported using onnxruntime-genai-cuda package.

[MadMenHitBooker]

@baijumeswani I installed the package using the command:
pip install --pre onnxruntime-genai-cuda --index-url=https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-genai/pypi/simple/