'update'

app.py

@@ -67,7 +67,7 @@ def sam_init():
     sam_checkpoint = os.path.join(os.path.dirname(__file__), "sam_pt", "sam_vit_h_4b8939.pth")
     model_type = "vit_h"
 
-    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device="cuda")
+    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device=f"cuda:{_GPU_ID}")
     predictor = SamPredictor(sam)
     return predictor
 
@@ -296,8 +296,11 @@ def run_demo():
 
     pipeline = load_era3d_pipeline(cfg)
     torch.set_grad_enabled(False)
-    pipeline.to('cuda')
+    print(1)
+    pipeline.to(device=f'cuda:{_GPU_ID}')
     pipeline.unet.enable_xformers_memory_efficient_attention()
+    
+    print(2)
     predictor = sam_init()