Skip to content

Fix --gpu_device mps falling back to CPU (fixes #1455)#1462

Open
gaoflow wants to merge 1 commit into
MouseLand:mainfrom
gaoflow:fix-1455-gpu-device-mps
Open

Fix --gpu_device mps falling back to CPU (fixes #1455)#1462
gaoflow wants to merge 1 commit into
MouseLand:mainfrom
gaoflow:fix-1455-gpu-device-mps

Conversation

@gaoflow
Copy link
Copy Markdown

@gaoflow gaoflow commented Jun 1, 2026

Summary

Fixes #1455. Running cellpose with --gpu_device mps on Apple silicon silently uses the CPU, while omitting the flag correctly uses MPS:

python -m cellpose --dir . --use_gpu --gpu_device mps ...
>>>> using CPU          # <- wrong, MPS is available

python -m cellpose --dir . --use_gpu ...
** TORCH MPS version installed and working. **
>>>> using GPU (MPS)    # <- works

Root cause

__main__ passes device=args.gpu_device (the string "mps") into assign_device, which keeps it as the string "mps" and calls use_gpu(gpu_number="mps"). _use_gpu_torch then treats gpu_number as a device index and builds:

torch.device("cuda:" + str(gpu_number))   # "cuda:mps" -> raises
torch.device("mps:"  + str(gpu_number))   # "mps:mps"  -> raises (invalid index)

Both raise, so _use_gpu_torch returns False and assign_device falls back to CPU. The default path works only because gpu_device defaults to "0", which becomes torch.device("mps:0").

Fix

In _use_gpu_torch, treat a non-integer gpu_number (i.e. the "mps" sentinel) as a request for the default MPS device torch.device("mps"), instead of building "mps:<gpu_number>". Integer indices keep their existing "mps:<n>" behaviour.

Verification

Reproduced on Apple silicon (MPS available):

from cellpose import core
core.assign_device(use_torch=True, gpu=True, device="mps")
# before: (device(type=cpu), False)
# after:  (device(type=mps), True)

Added a regression test test_assign_device_mps_string in tests/test_import.py (skipped when MPS is unavailable). It fails on the current code (assert gpu -> False) and passes with this change. The integer/default device paths are unchanged.

When the CLI is run with `--gpu_device mps`, the literal string "mps"
is passed through to `_use_gpu_torch` as the device index. It then
builds `torch.device("mps:mps")`, which is not a valid device string
and raises, so GPU detection fails and cellpose silently falls back to
CPU -- even though MPS is available (omitting the flag works).

Treat a non-integer gpu_number as a request for the default MPS device
(`torch.device("mps")`) instead of building `"mps:<gpu_number>"`.
Integer indices keep their existing `"mps:<n>"` behaviour.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Cellpose CLI uses CPU when --gpu_device mps is specified

1 participant