Merge pull request #5 from rbavery/cpu-mps-support

support mps device along with cpu and gpu
GeoCompas · Oct 11, 2024 · 94e7697 · 94e7697
2 parents db56c46 + a728fe5
commit 94e7697
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/app/requirements.txt b/app/requirements.txt
@@ -6,7 +6,7 @@ rtree==1.0.1
 typing-extensions>=4.10.0
 maplibre>=0.2.6
 shiny>=0.7.0
-segment-geospatial==0.11.0
+segment-geospatial==0.11.1
 geopandas==0.14.4
 shapely==2.0.4
 rasterio==1.3.9

diff --git a/app/utils/sam2.py b/app/utils/sam2.py
@@ -1,5 +1,5 @@
 import os
-from samgeo import SamGeo2
+from samgeo import SamGeo2, choose_device
 import torch
 from utils.utils import (
     generate_geojson,
@@ -10,7 +10,7 @@
 from utils.logger_config import log
 
 # Initialize the SAM model
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = choose_device()
 log.info(f"Using device: {device}")
 sam2 = SamGeo2(
     model_id="sam2-hiera-large",