logits to cpu

app.py

@@ -15,7 +15,7 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 print(f"Is CUDA available: {torch.cuda.is_available()} --> {device=}")
 if (torch.cuda.is_available()):
     print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
-    
+ 
 model.to(device)
 
 # https://github.com/NielsRogge/Transformers-Tutorials/blob/master/SegFormer/Segformer_inference_notebook.ipynb
@@ -69,22 +69,24 @@ def call(image): #nparray
     # resized_image = Image.fromarray(resized_image_np)
     # print(f"{resized_image=}")
 
-    inputs = feature_extractor(images=resized_image, return_tensors="pt")
-
+    inputs = feature_extractor(images=resized_image, return_tensors="pt").to(device)
+    
     print(f"**processing time: {(time.time() - start):.2f} s")
 
     outputs  = model(**inputs)
-    print(f"{outputs.logits.shape=}") # shape (batch_size, num_labels, height/4, width/4) -> 3, 19, 256 ,256
+    logits = outputs.logits.cpu()
+
+    print(f"{logits.shape=}") # shape (batch_size, num_labels, height/4, width/4) -> 3, 19, 256 ,256
     # print(f"{logits}")
 
     print(f"***processing time: {(time.time() - start):.2f} s")
     # First, rescale logits to original image size
     interpolated_logits =  nn.functional.interpolate(
-        outputs.logits,
+        logits,
         size=[1024, 1024], #resized_image.size[::-1], # (height, width)
         mode='bilinear',
         align_corners=False)
-    print(f"{interpolated_logits.shape=}, {outputs.logits.shape=}") # 1, 19, 1024, 1024
+    print(f"{interpolated_logits.shape=}, {logits.shape=}") # 1, 19, 1024, 1024
 
     # Second, apply argmax on the class dimension
     seg = interpolated_logits.argmax(dim=1)[0]