Allow server to run on specific GPU id

[che85]

It would be helpful to have an option to specify which GPU to use when running inference on a machine with multiple GPUs. In my case, I am running multiple MONAILabel servers, each with its own dedicated GPU.

[che85]

SlicerMONAIAuto3DSeg/MONAIAuto3DSeg/Scripts/auto3dseg_segresnet_inference.py

Line 83 in d748fd3

device = torch.device("cpu") if torch.cuda.device_count() == 0 else torch.device(0)

The script uses the first available GPU. We could use CUDA_VISIBILE_DEVICES environment variable, but I am not sure how this behaves when running a server which will run a subprocess with sys.executable

[lassoan]

Since we already pass several configuration parameters via command-line argument, we could just add one more optional device argument:

def main(model_file,
         image_file,
         result_file,
         save_mode=None,
         image_file_2=None,
         image_file_3=None,
         image_file_4=None,
         device=None,
         **kwargs):
...
if device is None:
    device = torch.device("cpu") if torch.cuda.device_count() == 0 else torch.device(0)
else:
    device = torch.device(device)
...