revert back to MPS after mask creation

Author: mrariden

cellpose/models.py

@@ -418,10 +418,11 @@ def _compute_masks(self, shape, dP, cellprob, flow_threshold=0.4, cellprob_thres
                        min_size=15, max_size_fraction=0.4, niter=None,
                        do_3D=False, stitch_threshold=0.0):
         """ compute masks from flows and cell probability """
-
+        changed_device_from = None
         if self.device.type == "mps" and do_3D:
             models_logger.warning("MPS does not support 3D post-processing, switching to CPU")
             self.device = torch.device("cpu")
+            changed_device_from = "mps"
         Lz, Ly, Lx = shape[:3]
         tic = time.time()
         if do_3D:
@@ -471,4 +472,7 @@ def _compute_masks(self, shape, dP, cellprob, flow_threshold=0.4, cellprob_thres
         if shape[0] > 1:
             models_logger.info("masks created in %2.2fs" % (flow_time))
 
+        if changed_device_from is not None:
+            models_logger.info("switching back to device %s" % self.device)
+            self.device = torch.device(changed_device_from)
         return masks