fix

app.py

@@ -72,8 +72,9 @@ def scan_checkpoints(ckpt_dir):
 
 def load_model(device,ckpt_dir='./checkpoints/mtlora.pt', pe_vision_config="PE-Core-L14-336"):
     """Load and configure model."""
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     backbone, transform, _ = get_backbone_pe(version='PE-Core-L14-336', apply_migration_flag=True, pretrained=False)
-    model = MTLModel(backbone,tasks=TASKS,use_lora=True,use_deep_head=True,
+    model = MTLModel(backbone,device=device,tasks=TASKS,use_lora=True,use_deep_head=True,
         use_mtl_lora=('mtlora' in ckpt_dir),
     )
     print(f'loading from {ckpt_dir}')
@@ -201,6 +202,7 @@ def init_model(ckpt_dir="./checkpoints/mtlora.pt", detection_confidence=0.5):
     # Load the perception encoder
     model, transform = load_model(ckpt_dir= ckpt_dir,device= device)
     model.eval()
+    print(device)
     model.to(device)