Update app.py

app.py

@@ -47,7 +47,8 @@ if 'confidence' not in st.session_state:
     st.session_state.confidence = 0.50
 
 if 'model' not in st.session_state:
-    st.session_state.model = YOLO("/home/user/app/models/head_model/best.pt", device='gpu')  # load 
+    torch.cuda.set_device(0)
+    st.session_state.model = YOLO("/home/user/app/models/head_model/best.pt")  # load 
 
 _, row0_col1, _ = st.columns([2,3,2])
 _, row1_col1, _ = st.columns([2,3,2])