NMS MPS device wrapper

[glenn-jocher]

May resolve #9613

Signed-off-by: Glenn Jocher glenn.jocher@ultralytics.com

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Improvement in MPS (Apple Silicon) support for Non-Maximum Suppression (NMS) in YOLOv5.

📊 Key Changes

• Added device detection to better handle Apple's Metal Performance Shaders (MPS).
• Introduced conditional processing to maintain tensor device consistency for MPS.
• Ensured that output tensors are moved back to the original device (if MPS) after NMS.

🎯 Purpose & Impact

• 🍏 Enhances compatibility with Apple Silicon by addressing MPS limitations.
• ⚡ Improves the efficiency of YOLOv5 on MPS-supported devices by ensuring proper tensor device handling.
• 💡 Allows developers and users with MPS-enabled devices to leverage native hardware acceleration without functionality issues in the NMS part of the inference.

[markusuntera]

There's still a problem. If some tensor's are empty (no boxes) these device is not set back to mps. Resulting multiple devices.

For example there should be before returning from "non_max_suppression":

    if mps:
        output = [a.to(device) for a in output]
    return output

[glenn-jocher]

@markusuntera thanks for sharing your feedback! It sounds like you are encountering an issue with the MPS device not being set back resulting in multiple devices when some tensors are empty. It seems like adding a conditional check to return the device to MPS if it is enabled before exiting the "non_max_suppression" function could be a viable solution. Have you tested this approach? Let me know if I can assist further!