Specify GPU Selection (e.g., CUDA:0, CUDA:1)

Hi,

Is there a way to specify which GPU to use for inference, such as restricting it to only `cuda:0` or `cuda:1` in the code? Or are there any workarounds for achieving this?

Thanks in advance.