Spaces:

henryscheible
/

barnacle-counter

Runtime error

Henry Scheible commited on Feb 22, 2023

Commit

a61a980

1 Parent(s): c3725ed

fix cuda error

Files changed (1) hide show

app.py CHANGED Viewed

@@ -50,7 +50,7 @@ def count_barnacles(input_img, progress=gr.Progress()):
     predicted_labels = torch.cat(predicted_labels_list)
     x = int(math.sqrt(predicted_labels.shape[0]))
     predicted_labels = predicted_labels.reshape([x, x, 2]).detach()
-    label_img = predicted_labels[:, :, :1].cuda().numpy()
     label_img -= label_img.min()
     label_img /= label_img.max()
     label_img = (label_img * 255).astype(np.uint8)

     predicted_labels = torch.cat(predicted_labels_list)
     x = int(math.sqrt(predicted_labels.shape[0]))
     predicted_labels = predicted_labels.reshape([x, x, 2]).detach()
+    label_img = predicted_labels[:, :, :1].cpu().numpy()
     label_img -= label_img.min()
     label_img /= label_img.max()
     label_img = (label_img * 255).astype(np.uint8)