From 017384aa6069a232ffe566e249f3b59606f43f26 Mon Sep 17 00:00:00 2001
From: GasTheWoods <124686986+Gia-Huynh@users.noreply.github.com>
Date: Mon, 14 Jul 2025 00:24:17 -0400
Subject: [PATCH] image2tensor select correct device

---
 depth_anything_v2/dpt.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/depth_anything_v2/dpt.py b/depth_anything_v2/dpt.py
index 18d3e6f8..9fde5427 100644
--- a/depth_anything_v2/dpt.py
+++ b/depth_anything_v2/dpt.py
@@ -214,8 +214,7 @@ def image2tensor(self, raw_image, input_size=518):
         
         image = transform({'image': image})['image']
         image = torch.from_numpy(image).unsqueeze(0)
-        
-        DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
+        DEVICE = next(self.pretrained.parameters()).device
         image = image.to(DEVICE)
         
         return image, (h, w)