From e4a12af62b126239a44581d265675081fb60403d Mon Sep 17 00:00:00 2001
From: Ryan Romanchuk <ryan@romanch.uk>
Date: Wed, 1 Apr 2026 15:25:30 -0500
Subject: [PATCH] Fix AIMV2 to use get_device() for MPS support

AIMV2 was hardcoded to select CUDA or CPU, skipping MPS. Use the
existing get_device() helper which already handles CUDA > MPS > CPU.
---
 classification.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/classification.py b/classification.py
index 5d518b2..7d5c481 100644
--- a/classification.py
+++ b/classification.py
@@ -418,7 +418,7 @@ def __init__(self, config):
         from transformers import AutoProcessor, AutoModel
 
         # Set up device
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = get_device()
 
         # Initialize model and processor
         self.processor = AutoProcessor.from_pretrained(