🗿 Killed bug that didn't let you go into cuda mode 🐛

SABS-R3-projects · Apr 15, 2020 · c54c61a · c54c61a
1 parent 43d5963
commit c54c61a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/Paint4Brains/Segmenter.py b/Paint4Brains/Segmenter.py
@@ -70,7 +70,7 @@ def _segment_over_one_axis(self, file_path, orientation):
             self.completion = self.completion + 50 / 256
             batch_x = volume[i:i + 1]
             if self.cuda_available and self.device == "cuda":
-                batch_x = batch_x.cuda(device)
+                batch_x = batch_x.cuda(self.device)
             volume_pred[i] = model(batch_x).cpu().numpy().astype(np.half)
 
         if orientation == "COR":