[BUG] MPS device not working with YOLO model

[nblumoe]

Describe the problem

FiftyOne's load_zoo_model supports passing in device="mps" to enable hardware acceleration on Apple Silicon.
While this is working with foz.load_zoo_model("resnet18-imagenet-torch", device="mps"), it doesn't work with a YOLO model: foz.load_zoo_model("yolov8x-world-torch", device="mps").

With the YOLO model, inference runs, but it isn't using the GPU acceleration.

Running YOLO manually and directly on the same machine is able to use GPU acceleration:

model_path = ...
model = YOLO(model_path)
batch_paths = ....
model(batch_paths, device=device)

Code to reproduce issue

import fiftyone.zoo as foz

dataset = foz.load_zoo_dataset("quickstart")
model = foz.load_zoo_model("yolov8x-world-torch", device="mps")

dataset.apply_model(
    model, 
    label_field="predictions_yoloworld",
)

Observe lacking GPU utilization in comparison to running YOLO on the GPU.

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 22.04): macOS 15.2 (24C101)
• Python version (python --version): Python 3.11.8
• FiftyOne version (fiftyone --version): FiftyOne v1.0.2, Voxel51, Inc.
• FiftyOne installed from (pip or source): pip / poetry

Other info/logs

I already tried os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1", but that wasn't sufficient to make it work.

Some additional validation on the MPS availability (but again, it also works when not running through FiftyOne, but with YOLO directly):

print(f"MPS available: {torch.backends.mps.is_available()}")
print(f"MPS built: {torch.backends.mps.is_built()}")
print(f"Default device: {torch.device(torch.empty(1).device)}")


MPS available: True
MPS built: True
Default device: mps:0

Willingness to contribute

The FiftyOne Community encourages bug fix contributions. Would you or another
member of your organization be willing to contribute a fix for this bug to the
FiftyOne codebase?

• Yes. I can contribute a fix for this bug independently
• Yes. I would be willing to contribute a fix for this bug with guidance
  from the FiftyOne community
• No. I cannot contribute a bug fix at this time

[nblumoe]

It looks like I can enable MPS acceleration like this:

from ultralytics import YOLO
import fiftyone.zoo as foz

model_path = ...

model = YOLO(model_path)
model = model.to("mps")

dataset = foz.load_zoo_dataset("quickstart")

dataset.apply_model(
    model, 
    label_field="predictions_yoloworld",
)