[BUG] Cellpose can only use gpu device 0
When I had 2 gpus available torch.cuda.device_count() == 2 and I tried to use gpu_device 1 when cuda:0 was already used - cellpose used the CPU instead of the GPU
The problem is that assign_device method never passes the device when it calls use_gpu(use_torch=True) and because of that use_gpu always tests device 0. So if GPU:0 is busy then assign_device will use CPU instead. The fix would be to call use_gpu(use_torch=use_torch,gpu_number=device)
Hi @cgoina, the core.assign_device() method does have a device argument that will select the CUDA device you want. This is exposed in the CLI (with the --gpu_device flag) but not exposed in the API.
Are you requesting that choosing your device is exposed in the API? Or can you tell me more about your issue?
The problem is that if you have 2 GPUs and GPU:0 is already in use if you call assign_device with GPU:1 it will give you the CPU not the GPU:1 even if GPU:1 is free. This happens because in assign_gpu when it calls use_gpu to test if the GPU is busy it does not forward the gpu_device parameter so use_gpu will always test only GPU:0 - I explained this to @carsen-stringer
@cgoina can you test PR #1311 for me?