Ooops, forgot no GPU on this instance

app.py

@@ -30,7 +30,7 @@ class DetectorModelOwl(nn.Module):
 		self.linear2 = nn.Linear(n_hidden * 2, 2)
 	
 	def forward(self, pixel_values: torch.Tensor, labels: torch.Tensor | None = None):
-		with torch.autocast("cuda", dtype=torch.bfloat16):
+		with torch.autocast("cpu", dtype=torch.bfloat16):
 			# Embed the image
 			outputs = self.owl(pixel_values=pixel_values, output_hidden_states=True)
 			x = outputs.last_hidden_state  # B, N, C
@@ -74,7 +74,7 @@ def owl_predict(image: Image.Image) -> bool:
 	input_image = TVF.normalize(preped, [0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711])
 
 	# Run
-	logits, = model(input_image.to('cuda').unsqueeze(0), None)
+	logits, = model(input_image.to('cpu').unsqueeze(0), None)
 	probs = F.softmax(logits, dim=1)
 	prediction = torch.argmax(probs.cpu(), dim=1)
 
@@ -106,7 +106,6 @@ def predict(image: Image.Image, conf_threshold: float):
 model = DetectorModelOwl("google/owlv2-base-patch16-ensemble", dropout=0.0)
 model.load_state_dict(torch.load("far5y1y5-8000.pt", map_location="cpu"))
 model.eval()
-model.cuda()
 
 # Load YOLO model
 yolo_model = YOLO("yolo11x-train28-best.pt")