fix a minor error

.idea/workspace.xml

@@ -5,8 +5,7 @@
   </component>
   <component name="ChangeListManager">
     <list default="true" id="9fb9e207-fc4f-4ff3-9adc-3c4c1e67daa7" name="Changes" comment="">
-      <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/utils.py" beforeDir="false" afterPath="$PROJECT_DIR$/utils.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/app.py" beforeDir="false" afterPath="$PROJECT_DIR$/app.py" afterDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
@@ -48,7 +47,7 @@
       <option name="presentableId" value="Default" />
       <updated>1664204268713</updated>
       <workItem from="1664204270261" duration="37000" />
-      <workItem from="1664204316867" duration="5840000" />
+      <workItem from="1664204316867" duration="6094000" />
     </task>
     <servers />
   </component>

app.py

@@ -54,7 +54,7 @@ def main(image: Image):
 
     # pred: H x W
     pred: torch.Tensor = logits.squeeze(dim=0).argmax(dim=0).cpu().numpy()
-    coloured_pred: np.ndarray = colourise_mask(mask=pred.cpu().numpy())
+    coloured_pred: np.ndarray = colourise_mask(mask=pred)
     super_imposed_img = cv2.addWeighted(coloured_pred, 0.5, np.array(pil_image), 0.5, 0)
 
     # resize prediction to original resolution