Repository: juglab/EmbedSeg (language: python)

Issue:
Possibility to train on non-GPU machines?
Hello team, I am trying to run the 3 DSB notebooks from the examples on my macbook (no GPU), and I run into the error `FileNotFoundError: [Errno 2] No such file or directory: 'nvidia-smi': 'nvidia-smi'`. This part of the code exists in the `EmbedSeg/utils/preprocess_data.py` location, specifically `def get_gpu_memory()`.
Is there a way to "switch-off" this get gpu memory thing for training on non-gpu machines?

Thanks, Ajinkya 



Discussion / hints:
Hi team, any update on this?
Hi @ajinkya-kulkarni,

Sorry this repository is not really maintained, but we have plans to revamp it early next year. 

In the mean time, I have merged your contributions. Thanks a lot!  I guess there the `get_gpu_memory()` problem left, we will try to have a look.

Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..f0b53ad
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,22 @@
+ci:
+  autoupdate_schedule: monthly
+  autofix_commit_msg: "style(pre-commit.ci): auto fixes [...]"
+  autoupdate_commit_msg: "ci(pre-commit.ci): autoupdate"
+
+repos:
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
+    rev: v0.0.254
+    hooks:
+      - id: ruff
+        args: [--fix]
+
+  - repo: https://github.com/psf/black
+    rev: 23.1.0
+    hooks:
+      - id: black
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.1.1
+    hooks:
+      - id: mypy
+        files: "^EmbedSeg/"
diff --git a/EmbedSeg/criterions/__init__.py b/EmbedSeg/criterions/__init__.py
index b11cec2..fa528be 100644
--- a/EmbedSeg/criterions/__init__.py
+++ b/EmbedSeg/criterions/__init__.py
@@ -2,10 +2,19 @@
 from EmbedSeg.criterions.my_loss_3d import SpatialEmbLoss_3d
 
 
-
-def get_loss(grid_z, grid_y, grid_x, pixel_z, pixel_y, pixel_x, one_hot, loss_opts):
+def get_loss(
+    grid_z,
+    grid_y,
+    grid_x,
+    pixel_z,
+    pixel_y,
+    pixel_x,
+    one_hot,
+    loss_opts,
+):
     if grid_z is None:
         return SpatialEmbLoss(grid_y, grid_x, pixel_y, pixel_x, one_hot, **loss_opts)
     else:
-        return SpatialEmbLoss_3d(grid_z, grid_y, grid_x, pixel_z, pixel_y, pixel_x, one_hot, **loss_opts)
-    
\ No newline at end of file
+        return SpatialEmbLoss_3d(
+            grid_z, grid_y, grid_x, pixel_z, pixel_y, pixel_x, one_hot, **loss_opts
+        )
diff --git a/EmbedSeg/criterions/lovasz_losses.py b/EmbedSeg/criterions/lovasz_losses.py
index 5ee5c68..31db3c0 100644
--- a/EmbedSeg/criterions/lovasz_losses.py
+++ b/EmbedSeg/criterions/lovasz_losses.py
@@ -9,10 +9,11 @@
 from torch.autograd import Variable
 import torch.nn.functional as F
 import numpy as np
+
 try:
     from itertools import ifilterfalse
 except ImportError:  # py3k
-    from itertools import filterfalse
+    pass
 
 
 def lovasz_grad(gt_sorted):
@@ -23,15 +24,15 @@ def lovasz_grad(gt_sorted):
     p = len(gt_sorted)
     gts = gt_sorted.sum()
     intersection = gts.float() - gt_sorted.float().cumsum(0)
-    #union = gts.float() + (1 - gt_sorted).float().cumsum(0)
+    # union = gts.float() + (1 - gt_sorted).float().cumsum(0)
     union = gts.float() + (1 - gt_sorted.float()).cumsum(0)
-    jaccard = 1. - intersection / union
+    jaccard = 1.0 - intersection / union
     if p > 1:  # cover 1-pixel case
         jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
     return jaccard
 
 
-def iou_binary(preds, labels, EMPTY=1., ignore=None, per_image=True):
+def iou_binary(preds, labels, EMPTY=1.0, ignore=None, per_image=True):
     """
     IoU for foreground class
     binary: 1 foreground, 0 background
@@ -47,11 +48,11 @@ def iou_binary(preds, labels, EMPTY=1., ignore=None, per_image=True):
         else:
             iou = float(intersection) / union
         ious.append(iou)
-    iou = mean(ious)    # mean accross images if per_image
+    iou = mean(ious)  # mean accross images if per_image
     return 100 * iou
 
 
-def iou(preds, labels, C, EMPTY=1., ignore=None, per_image=False):
+def iou(preds, labels, C, EMPTY=1.0, ignore=None, per_image=False):
     """
     Array of IoU for each (non ignored) class
     """
@@ -64,8 +65,7 @@ def iou(preds, labels, C, EMPTY=1., ignore=None, per_image=False):
             # The ignored label is sometimes among predicted classes (ENet - CityScapes)
             if i != ignore:
                 intersection = ((label == i) & (pred == i)).sum()
-                union = ((label == i) | (
-                    (pred == i) & (label != ignore))).sum()
+                union = ((label == i) | ((pred == i) & (label != ignore))).sum()
                 if not union:
                     iou.append(EMPTY)
                 else:
@@ -87,11 +87,14 @@ def lovasz_hinge(logits, labels, per_image=True, ignore=None):
       ignore: void class id
     """
     if per_image:
-        loss = mean(lovasz_hinge_flat(*flatten_binary_scores(log.unsqueeze(0), lab.unsqueeze(0), ignore))
-                    for log, lab in zip(logits, labels))
+        loss = mean(
+            lovasz_hinge_flat(
+                *flatten_binary_scores(log.unsqueeze(0), lab.unsqueeze(0), ignore)
+            )
+            for log, lab in zip(logits, labels)
+        )
     else:
-        loss = lovasz_hinge_flat(
-            *flatten_binary_scores(logits, labels, ignore))
+        loss = lovasz_hinge_flat(*flatten_binary_scores(logits, labels, ignore))
     return loss
 
 
@@ -104,9 +107,9 @@ def lovasz_hinge_flat(logits, labels):
     """
     if len(labels) == 0:
         # only void pixels, the gradients should be 0
-        return logits.sum() * 0.
-    signs = 2. * labels.float() - 1.
-    errors = (1. - logits * Variable(signs))
+        return logits.sum() * 0.0
+    signs = 2.0 * labels.float() - 1.0
+    errors = 1.0 - logits * Variable(signs)
     errors_sorted, perm = torch.sort(errors, dim=0, descending=True)
     perm = perm.data
     gt_sorted = labels[perm]
@@ -124,7 +127,7 @@ def flatten_binary_scores(scores, labels, ignore=None):
     labels = labels.view(-1)
     if ignore is None:
         return scores, labels
-    valid = (labels != ignore)
+    valid = labels != ignore
     vscores = scores[valid]
     vlabels = labels[valid]
     return vscores, vlabels
@@ -135,7 +138,7 @@ def __init__(self):
         super(StableBCELoss, self).__init__()
 
     def forward(self, input, target):
-        neg_abs = - input.abs()
+        neg_abs = -input.abs()
         loss = input.clamp(min=0) - input * target + (1 + neg_abs.exp()).log()
         return loss.mean()
 
@@ -158,26 +161,36 @@ def binary_xloss(logits, labels, ignore=None):
 def lovasz_softmax(probas, labels, only_present=False, per_image=False, ignore=None):
     """
     Multi-class Lovasz-Softmax loss
-      probas: [B, C, H, W] Variable, class probabilities at each prediction (between 0 and 1)
-      labels: [B, H, W] Tensor, ground truth labels (between 0 and C - 1)
+      probas: [B, C, H, W] Variable,
+            class probabilities at each prediction (between 0 and 1)
+      labels: [B, H, W] Tensor,
+            ground truth labels (between 0 and C - 1)
       only_present: average only on classes present in ground truth
       per_image: compute the loss per image instead of per batch
       ignore: void class labels
     """
     if per_image:
-        loss = mean(lovasz_softmax_flat(*flatten_probas(prob.unsqueeze(0), lab.unsqueeze(0), ignore), only_present=only_present)
-                    for prob, lab in zip(probas, labels))
+        loss = mean(
+            lovasz_softmax_flat(
+                *flatten_probas(prob.unsqueeze(0), lab.unsqueeze(0), ignore),
+                only_present=only_present
+            )
+            for prob, lab in zip(probas, labels)
+        )
     else:
         loss = lovasz_softmax_flat(
-            *flatten_probas(probas, labels, ignore), only_present=only_present)
+            *flatten_probas(probas, labels, ignore), only_present=only_present
+        )
     return loss
 
 
 def lovasz_softmax_flat(probas, labels, only_present=False):
     """
     Multi-class Lovasz-Softmax loss
-      probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1)
-      labels: [P] Tensor, ground truth labels (between 0 and C - 1)
+      probas: [P, C] Variable,
+            class probabilities at each prediction (between 0 and 1)
+      labels: [P] Tensor,
+            ground truth labels (between 0 and C - 1)
       only_present: average only on classes present in ground truth
     """
     C = probas.size(1)
@@ -190,8 +203,7 @@ def lovasz_softmax_flat(probas, labels, only_present=False):
         errors_sorted, perm = torch.sort(errors, 0, descending=True)
         perm = perm.data
         fg_sorted = fg[perm]
-        losses.append(
-            torch.dot(errors_sorted, Variable(lovasz_grad(fg_sorted))))
+        losses.append(torch.dot(errors_sorted, Variable(lovasz_grad(fg_sorted))))
     return mean(losses)
 
 
@@ -200,12 +212,11 @@ def flatten_probas(probas, labels, ignore=None):
     Flattens predictions in the batch
     """
     B, C, H, W = probas.size()
-    probas = probas.permute(0, 2, 3, 1).contiguous(
-    ).view(-1, C)  # B * H * W, C = P, C
+    probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C)  # B * H * W, C = P, C
     labels = labels.view(-1)
     if ignore is None:
         return probas, labels
-    valid = (labels != ignore)
+    valid = labels != ignore
     vprobas = probas[valid.nonzero().squeeze()]
     vlabels = labels[valid]
     return vprobas, vlabels
@@ -220,6 +231,7 @@ def xloss(logits, labels, ignore=None):
 
 # --------------------------- HELPER FUNCTIONS ---------------------------
 
+
 def mean(l, ignore_nan=False, empty=0):
     """
     nanmean compatible with generators.
@@ -231,8 +243,8 @@ def mean(l, ignore_nan=False, empty=0):
         n = 1
         acc = next(l)
     except StopIteration:
-        if empty == 'raise':
-            raise ValueError('Empty mean')
+        if empty == "raise":
+            raise ValueError("Empty mean")
         return empty
     for n, v in enumerate(l, 2):
         acc += v
diff --git a/EmbedSeg/criterions/my_loss.py b/EmbedSeg/criterions/my_loss.py
index e2f3974..e026c47 100644
--- a/EmbedSeg/criterions/my_loss.py
+++ b/EmbedSeg/criterions/my_loss.py
@@ -6,51 +6,74 @@
 
 
 class SpatialEmbLoss(nn.Module):
-
-    def __init__(self, grid_y=1024, grid_x=1024, pixel_y=1, pixel_x=1, one_hot=False, n_sigma=2, foreground_weight=1):
+    def __init__(
+        self,
+        grid_y=1024,
+        grid_x=1024,
+        pixel_y=1,
+        pixel_x=1,
+        one_hot=False,
+        n_sigma=2,
+        foreground_weight=1,
+    ):
         super().__init__()
 
-        print('Created spatial emb loss function with: n_sigma: {}, foreground_weight: {}'.format(
-            n_sigma, foreground_weight))
-        print('*************************')
+        print(
+            "Created spatial emb loss function with: \
+                    n_sigma: {}, foreground_weight: {}".format(
+                n_sigma, foreground_weight
+            )
+        )
+        print("*************************")
         self.n_sigma = n_sigma
         self.foreground_weight = foreground_weight
 
         # coordinate map
-        xm = torch.linspace(0, pixel_x, grid_x).view(
-            1, 1, -1).expand(1, grid_y, grid_x)
-        ym = torch.linspace(0, pixel_y, grid_y).view(
-            1, -1, 1).expand(1, grid_y, grid_x)
+        xm = torch.linspace(0, pixel_x, grid_x).view(1, 1, -1).expand(1, grid_y, grid_x)
+        ym = torch.linspace(0, pixel_y, grid_y).view(1, -1, 1).expand(1, grid_y, grid_x)
         xym = torch.cat((xm, ym), 0)
-
         self.register_buffer("xym", xym)
         self.one_hot = one_hot
 
-    def forward(self, prediction, instances, labels, center_images, w_inst=1, w_var=10, w_seed=1, iou=False,
-                iou_meter=None):
-
+    def forward(
+        self,
+        prediction,
+        instances,
+        labels,
+        center_images,
+        w_inst=1,
+        w_var=10,
+        w_seed=1,
+        iou=False,
+        iou_meter=None,
+    ):
         # instances B 1 Y X
-        batch_size, height, width = prediction.size(0), prediction.size(2), prediction.size(3)
+        batch_size, height, width = (
+            prediction.size(0),
+            prediction.size(2),
+            prediction.size(3),
+        )
 
         xym_s = self.xym[:, 0:height, 0:width].contiguous()  # 2 x h x w
 
         loss = 0
 
         for b in range(0, batch_size):
-
             spatial_emb = torch.tanh(prediction[b, 0:2]) + xym_s  # 2 x h x w #TODO
-            sigma = prediction[b, 2:2 + self.n_sigma]  # n_sigma x h x w
-            seed_map = torch.sigmoid(prediction[b, 2 + self.n_sigma:2 + self.n_sigma + 1])  # 1 x h x w
+            sigma = prediction[b, 2 : 2 + self.n_sigma]  # n_sigma x h x w
+            seed_map = torch.sigmoid(
+                prediction[b, 2 + self.n_sigma : 2 + self.n_sigma + 1]
+            )  # 1 x h x w
             # loss accumulators
             var_loss = 0
             instance_loss = 0
             seed_loss = 0
             obj_count = 0
 
-            instance = instances[b].unsqueeze(0)  # 1 x h x w or 1 x d x h x w (one_hot)  
-            label = labels[b].unsqueeze(0)  # 1 x h x w 
+            instance = instances[b].unsqueeze(0)  # 1 x h x w or 1 x d x h x w (one_hot)
+            label = labels[b].unsqueeze(0)  # 1 x h x w
             center_image = center_images[b].unsqueeze(0)  # 1 x h x w
-            if (self.one_hot):
+            if self.one_hot:
                 instance_ids = np.arange(instances.size(1))
             else:
                 instance_ids = instance.unique()
@@ -63,16 +86,18 @@ def forward(self, prediction, instances, labels, center_images, w_inst=1, w_var=
                 seed_loss += torch.sum(torch.pow(seed_map[bg_mask] - 0, 2))
 
             for id in instance_ids:
-                if (self.one_hot):
+                if self.one_hot:
                     in_mask = instance[:, id, ...].eq(1)
                 else:
                     in_mask = instance.eq(id)  # 1 x h x w
-                #center_mask = in_mask & center_image.byte()
+                # center_mask = in_mask & center_image.byte()
                 center_mask = in_mask & center_image
-                if (center_mask.sum().eq(1)):
+                if center_mask.sum().eq(1):
                     center = xym_s[center_mask.expand_as(xym_s)].view(2, 1, 1)
                 else:
-                    xy_in = xym_s[in_mask.expand_as(xym_s)].view(2, -1)  # TODO --> should this edge case change!
+                    xy_in = xym_s[in_mask.expand_as(xym_s)].view(
+                        2, -1
+                    )  # TODO --> should this edge case change!
                     center = xy_in.mean(1).view(2, 1, 1)  # 2 x 1 x 1
 
                 # calculate sigma
@@ -81,20 +106,22 @@ def forward(self, prediction, instances, labels, center_images, w_inst=1, w_var=
                 s = sigma_in.mean(1).view(self.n_sigma, 1, 1)  # n_sigma x 1 x 1
 
                 # calculate var loss before exp
-                var_loss = var_loss + \
-                           torch.mean(
-                               torch.pow(sigma_in - s[..., 0].detach(), 2))
+                var_loss = var_loss + torch.mean(
+                    torch.pow(sigma_in - s[..., 0].detach(), 2)
+                )
 
                 s = torch.exp(s * 10)  # TODO
-                dist = torch.exp(-1 * torch.sum(torch.pow(spatial_emb - center, 2) * s, 0, keepdim=True))
+                dist = torch.exp(
+                    -1
+                    * torch.sum(torch.pow(spatial_emb - center, 2) * s, 0, keepdim=True)
+                )
 
                 # apply lovasz-hinge loss
-                instance_loss = instance_loss + \
-                                lovasz_hinge(dist * 2 - 1, in_mask)
-
+                instance_loss = instance_loss + lovasz_hinge(dist * 2 - 1, in_mask)
                 # seed loss
                 seed_loss += self.foreground_weight * torch.sum(
-                    torch.pow(seed_map[in_mask] - dist[in_mask].detach(), 2))
+                    torch.pow(seed_map[in_mask] - dist[in_mask].detach(), 2)
+                )
 
                 # calculate instance iou
                 if iou:
diff --git a/EmbedSeg/criterions/my_loss_3d.py b/EmbedSeg/criterions/my_loss_3d.py
index bd9a8d6..4053670 100644
--- a/EmbedSeg/criterions/my_loss_3d.py
+++ b/EmbedSeg/criterions/my_loss_3d.py
@@ -1,43 +1,85 @@
-import numpy as np
 import torch
 import torch.nn as nn
 from EmbedSeg.criterions.lovasz_losses import lovasz_hinge
 
 
 class SpatialEmbLoss_3d(nn.Module):
-
-    def __init__(self, grid_z = 32, grid_y=1024, grid_x=1024, pixel_z = 1, pixel_y=1, pixel_x=1, one_hot=False, n_sigma=3, foreground_weight=1):
+    def __init__(
+        self,
+        grid_z=32,
+        grid_y=1024,
+        grid_x=1024,
+        pixel_z=1,
+        pixel_y=1,
+        pixel_x=1,
+        one_hot=False,
+        n_sigma=3,
+        foreground_weight=1,
+    ):
         super().__init__()
 
-        print('Created spatial emb loss function with: n_sigma: {}, foreground_weight: {}'.format(
-            n_sigma, foreground_weight))
-        print('*************************')
+        print(
+            "Created spatial emb loss function with: \
+                    n_sigma: {}, foreground_weight: {}".format(
+                n_sigma, foreground_weight
+            )
+        )
+        print("*************************")
         self.n_sigma = n_sigma
         self.foreground_weight = foreground_weight
 
-        xm = torch.linspace(0, pixel_x, grid_x).view(1, 1, 1, -1).expand(1, grid_z, grid_y, grid_x)
-        ym = torch.linspace(0, pixel_y, grid_y).view(1, 1, -1, 1).expand(1, grid_z, grid_y, grid_x)
-        zm = torch.linspace(0, pixel_z, grid_z).view(1, -1, 1, 1).expand(1, grid_z, grid_y, grid_x)
+        xm = (
+            torch.linspace(0, pixel_x, grid_x)
+            .view(1, 1, 1, -1)
+            .expand(1, grid_z, grid_y, grid_x)
+        )
+        ym = (
+            torch.linspace(0, pixel_y, grid_y)
+            .view(1, 1, -1, 1)
+            .expand(1, grid_z, grid_y, grid_x)
+        )
+        zm = (
+            torch.linspace(0, pixel_z, grid_z)
+            .view(1, -1, 1, 1)
+            .expand(1, grid_z, grid_y, grid_x)
+        )
         xyzm = torch.cat((xm, ym, zm), 0)
 
         self.register_buffer("xyzm", xyzm)
         self.one_hot = one_hot
 
-    def forward(self, prediction, instances, labels, center_images, w_inst=1, w_var=10, w_seed=1, iou=False,
-                iou_meter=None):
-
+    def forward(
+        self,
+        prediction,
+        instances,
+        labels,
+        center_images,
+        w_inst=1,
+        w_var=10,
+        w_seed=1,
+        iou=False,
+        iou_meter=None,
+    ):
         # instances B 1 Z Y X
-        batch_size, depth, height, width = prediction.size(0), prediction.size(2), prediction.size(3), prediction.size(4)
+        batch_size, depth, height, width = (
+            prediction.size(0),
+            prediction.size(2),
+            prediction.size(3),
+            prediction.size(4),
+        )
 
         xyzm_s = self.xyzm[:, 0:depth, 0:height, 0:width].contiguous()  # 3 x d x h x w
 
+        print(f"xyzm_s shape is {xyzm_s.shape}")
+        print(f"prediction shape is {prediction.shape}")
         loss = 0
 
         for b in range(0, batch_size):
-
-            spatial_emb = torch.tanh(prediction[b, 0:3]) + xyzm_s  # 3 x d x h x w 
-            sigma = prediction[b, 3:3 + self.n_sigma]  # n_sigma x d x h x w
-            seed_map = torch.sigmoid(prediction[b, 3 + self.n_sigma:3 + self.n_sigma + 1])  # 1 x d x h x w 
+            spatial_emb = torch.tanh(prediction[b, 0:3]) + xyzm_s  # 3 x d x h x w
+            sigma = prediction[b, 3 : 3 + self.n_sigma]  # n_sigma x d x h x w
+            seed_map = torch.sigmoid(
+                prediction[b, 3 + self.n_sigma : 3 + self.n_sigma + 1]
+            )  # 1 x d x h x w
             # loss accumulators
             var_loss = 0
             instance_loss = 0
@@ -57,34 +99,40 @@ def forward(self, prediction, instances, labels, center_images, w_inst=1, w_var=
                 seed_loss += torch.sum(torch.pow(seed_map[bg_mask] - 0, 2))
 
             for id in instance_ids:
-
                 in_mask = instance.eq(id)  # 1 x d x h x w
-                #center_mask = in_mask & center_image.byte()
+                # center_mask = in_mask & center_image.byte()
                 center_mask = in_mask & center_image
-                if (center_mask.sum().eq(1)):
+                if center_mask.sum().eq(1):
                     center = xyzm_s[center_mask.expand_as(xyzm_s)].view(3, 1, 1, 1)
                 else:
-                    xyz_in = xyzm_s[in_mask.expand_as(xyzm_s)].view(3, -1)  
+                    xyz_in = xyzm_s[in_mask.expand_as(xyzm_s)].view(3, -1)
                     center = xyz_in.mean(1).view(3, 1, 1, 1)  # 3 x 1 x 1 x 1
 
                 # calculate sigma
-                sigma_in = sigma[in_mask.expand_as(sigma)].view(self.n_sigma, -1) # 3 x N
+                sigma_in = sigma[in_mask.expand_as(sigma)].view(
+                    self.n_sigma, -1
+                )  # 3 x N
 
                 s = sigma_in.mean(1).view(self.n_sigma, 1, 1, 1)  # n_sigma x 1 x 1 x 1
 
                 # calculate var loss before exp
-                var_loss = var_loss + torch.mean( torch.pow(sigma_in - s[..., 0, 0].detach(), 2))
+                var_loss = var_loss + torch.mean(
+                    torch.pow(sigma_in - s[..., 0, 0].detach(), 2)
+                )
 
                 s = torch.exp(s * 10)
-                dist = torch.exp(-1 * torch.sum(torch.pow(spatial_emb - center, 2) * s, 0, keepdim=True))
+                dist = torch.exp(
+                    -1
+                    * torch.sum(torch.pow(spatial_emb - center, 2) * s, 0, keepdim=True)
+                )
 
                 # apply lovasz-hinge loss
-                instance_loss = instance_loss + \
-                                lovasz_hinge(dist * 2 - 1, in_mask)
+                instance_loss = instance_loss + lovasz_hinge(dist * 2 - 1, in_mask)
 
                 # seed loss
                 seed_loss += self.foreground_weight * torch.sum(
-                    torch.pow(seed_map[in_mask] - dist[in_mask].detach(), 2))
+                    torch.pow(seed_map[in_mask] - dist[in_mask].detach(), 2)
+                )
 
                 # calculate instance iou
                 if iou:
diff --git a/EmbedSeg/datasets/ThreeDimensionalDataset.py b/EmbedSeg/datasets/ThreeDimensionalDataset.py
index 0fe0e65..9443220 100644
--- a/EmbedSeg/datasets/ThreeDimensionalDataset.py
+++ b/EmbedSeg/datasets/ThreeDimensionalDataset.py
@@ -10,94 +10,142 @@
 from skimage.segmentation import relabel_sequential
 from torch.utils.data import Dataset
 
-from EmbedSeg.utils.generate_crops import normalize_min_max_percentile, normalize_mean_std
+from EmbedSeg.utils.generate_crops import (
+    normalize_min_max_percentile,
+    normalize_mean_std,
+)
 
 
 class ThreeDimensionalDataset(Dataset):
     """
-            A class used to create a PyTorch Dataset for handling 3D images and label masks
+    A class used to create a PyTorch Dataset for handling 3D
+    images and label masks
 
-            Attributes
-            ----------
-            image_list: list of strings containing paths to the images
+    Attributes
+    ----------
+    image_list: list of strings containing paths to the images
 
-            instance_list: list of strings containing paths to the GT instances
+    instance_list: list of strings containing paths to the GT instances
 
-            center_image_list: list of strings containing paths to the center images
+    center_image_list: list of strings containing paths to the center images
 
-            bg_id: int
-                Id corresponding to the background, default=0
-            size: int
-                This is set equal to a different value than `real_size` in case the epoch is desired to be shorter
+    bg_id: int
+        Id corresponding to the background, default=0
 
-            real_size: int
-                Actual number of the number of images
+    size: int
+        This is set equal to a different value than `real_size`
+        in case the epoch is desired to be shorter
 
-            transform: PyTorch transform
+    real_size: int
+        Actual number of the number of images
 
-            one_hot: bool, default = False
-                Should be set equal to True if the GT masks are available in a one-hot encoded fashion
-                This is not applicable in a 3D setting since any pixel can only be assigned to one 3D GT object
-                This parameter will be deprecated for the 3D setting in a future release
+    transform: PyTorch transform
 
-            norm: str
-                Should be set equal to one of `min-max-percentile`, `absolute`, `mean-std`
+    one_hot: bool, default = False
+        Should be set equal to True if the GT masks are available
+        in a one-hot encoded fashion
+        This is not applicable in a 3D setting since any pixel
+        can only be assigned to one 3D GT object
+        This parameter will be deprecated for the 3D setting in a future release
 
-            type: str
-                Should be set equal to one of `train`, `val` or `test`
-            data_type: str
-                Should be set equal to one of `8-bit` or `16-bit`
+    norm: str
+        Should be set equal to one of `min-max-percentile`, `absolute`, `mean-std`
 
-            normalization: bool
-                Should be equal to True for test images and set equal to False during training (since crops are already normalized)
-            anisotropy_factor: float
-                Should be set equal to the ratio of the pixel sizes along the z dimension and the x dimension
-                In case, the raw image is acquired such that the image is down-sampled along the z-diemnsion,
-                then 'anisotropy_factor' should be greater than equal to 1.0
+    type: str
+        Should be set equal to one of `train`, `val` or `test`
+    data_type: str
+        Should be set equal to one of `8-bit` or `16-bit`
 
-            sliced_mode: bool
-                In case the 3D images should be interpreted in a 'sliced fashion' (i.e. a 2D model should be used for training on individual slices
-                And the predictions from these slices considered independently are combined during inference to generate a 3D instance mask), then
-                'sliced_mode' should be set equal to True
+    normalization: bool
+        Should be equal to True for test images and set equal to False
+        during training (since crops are already normalized)
+    anisotropy_factor: float
+        Should be set equal to the ratio of the pixel sizes
+        along the z dimension and the x dimension
+        In case, the raw image is acquired such that the image
+        is down-sampled along the z-diemnsion,
+        then 'anisotropy_factor' should be greater than equal to 1.0
 
-            uniform_ds_factor: int, default = 1
-                If the original crops were generated by down-sampling in '01-data.ipynb', then the test images
-                should also be down-sampled before being input to the model and the predictions should then be up-sampled
-                to restore them back to the original size
+    sliced_mode: bool
+        In case the 3D images should be interpreted in a 'sliced fashion'
+        (i.e. a 2D model should be used for training on individual slices
+        And the predictions from these slices considered independently are
+         combined during inference to generate a 3D instance mask), then
+        'sliced_mode' should be set equal to True
 
-            Methods
-            -------
-            __init__: Initializes an object of class `ThreeDimensionalDataset`
+    uniform_ds_factor: int, default = 1
+        If the original crops were generated by down-sampling
+        in '01-data.ipynb', then the test images
+        should also be down-sampled before being input to the model
+        and the predictions should then be up-sampled
+        to restore them back to the original size
 
-            __len__: Returns self.real_size if self.size = None
+    Methods
+    -------
+    __init__: Initializes an object of class `ThreeDimensionalDataset`
 
-            convert_yx_to_cyx: Adds an additional dimension for channel
+    __len__: Returns self.real_size if self.size = None
 
-            decode_instance: In case, instances are available as tiffs, this method decodes them
+    convert_yx_to_cyx: Adds an additional dimension for channel
 
-            rle_decode: In case, instances are available as csv files (and not tiffs), this method decodes them
-            """
+    decode_instance: In case, instances are available as tiffs,
+        this method decodes them
 
-    def __init__(self, data_dir='./', center='center-medoid', type="train", bg_id=0, size=None, transform=None,
-                 one_hot=False, norm='min-max-percentile', normalization=False, data_type='8-bit',
-                 anisotropy_factor=1.0, sliced_mode=False, uniform_ds_factor = 1):
-        print('3-D `{}` dataloader created! Accessing data from {}/{}/'.format(type, data_dir, type))
+    rle_decode: In case, instances are available as csv files (and not tiffs),
+        this method decodes them
+    """
+
+    def __init__(
+        self,
+        data_dir="./",
+        center="center-medoid",
+        type="train",
+        bg_id=0,
+        size=None,
+        transform=None,
+        one_hot=False,
+        norm="min-max-percentile",
+        normalization=False,
+        data_type="8-bit",
+        anisotropy_factor=1.0,
+        sliced_mode=False,
+        uniform_ds_factor=1,
+    ):
+        print(
+            "3-D `{}` dataloader created! Accessing data from {}/{}/".format(
+                type, data_dir, type
+            )
+        )
 
         # get image and instance list
-        image_list = glob.glob(os.path.join(data_dir, '{}/'.format(type), 'images/*.tif'))
+        image_list = glob.glob(
+            os.path.join(data_dir, "{}/".format(type), "images/*.tif")
+        )
         image_list.sort()
-        print('Number of images in `{}` directory is {}'.format(type, len(image_list)))
+        print("Number of images in `{}` directory is {}".format(type, len(image_list)))
         self.image_list = image_list
 
-        instance_list = glob.glob(os.path.join(data_dir, '{}/'.format(type), 'masks/*.tif'))
+        instance_list = glob.glob(
+            os.path.join(data_dir, "{}/".format(type), "masks/*.tif")
+        )
         instance_list.sort()
-        print('Number of instances in `{}` directory is {}'.format(type, len(instance_list)))
+        print(
+            "Number of instances in `{}` directory is {}".format(
+                type, len(instance_list)
+            )
+        )
         self.instance_list = instance_list
 
-        center_image_list = glob.glob(os.path.join(data_dir, '{}/'.format(type), center + '/*.tif'))
+        center_image_list = glob.glob(
+            os.path.join(data_dir, "{}/".format(type), center + "/*.tif")
+        )
         center_image_list.sort()
-        print('Number of center images in `{}` directory is {}'.format(type, len(center_image_list)))
-        print('*************************')
+        print(
+            "Number of center images in `{}` directory is {}".format(
+                type, len(center_image_list)
+            )
+        )
+        print("*************************")
         self.center_image_list = center_image_list
 
         self.bg_id = bg_id
@@ -126,49 +174,61 @@ def __getitem__(self, index):
 
         # load image
         image = tifffile.imread(self.image_list[index])  # ZYX
-        if (self.normalization and self.norm == 'min-max-percentile'):
+        if self.normalization and self.norm == "min-max-percentile":
             image = normalize_min_max_percentile(image, 1, 99.8, axis=(0, 1, 2))
-        elif (self.normalization and self.norm == 'mean-std'):
+        elif self.normalization and self.norm == "mean-std":
             image = normalize_mean_std(image)
-        elif (self.normalization and self.norm == 'absolute'):
-            image = image.astype(np.float32)
-            if self.data_type == '8-bit':
+        elif self.normalization and self.norm == "absolute":
+            image = image.astype(float)
+            if self.data_type == "8-bit":
                 image /= 255
-            elif self.data_type == '16-bit':
+            elif self.data_type == "16-bit":
                 image /= 65535
-        if self.type == 'test' and self.sliced_mode:
+        if self.type == "test" and self.sliced_mode:
             image = zoom(image, (self.anisotropy_factor, 1, 1), order=0)
-        elif self.type == 'test' and not self.sliced_mode:
-            image = image[::self.uniform_ds_factor, ::self.uniform_ds_factor, ::self.uniform_ds_factor]
-        image = self.convert_zyx_to_czyx(image, key='image')  # CZYX
-        sample['image'] = image  # CZYX
-        sample['im_name'] = self.image_list[index]
-        if (len(self.instance_list) != 0):
-            if self.instance_list[index][-3:] == 'csv':
-                instance = self.rle_decode(self.instance_list[index], one_hot=self.one_hot)  # ZYX
+        elif self.type == "test" and not self.sliced_mode:
+            image = image[
+                :: self.uniform_ds_factor,
+                :: self.uniform_ds_factor,
+                :: self.uniform_ds_factor,
+            ]
+        image = self.convert_zyx_to_czyx(image, key="image")  # CZYX
+        sample["image"] = image  # CZYX
+        sample["im_name"] = self.image_list[index]
+        if len(self.instance_list) != 0:
+            if self.instance_list[index][-3:] == "csv":
+                instance = self.rle_decode(
+                    self.instance_list[index], one_hot=self.one_hot
+                )  # ZYX
             else:
                 instance = tifffile.imread(self.instance_list[index])  # ZYX
 
             instance, label = self.decode_instance(instance, self.one_hot, self.bg_id)
-            if self.type == 'test' and self.sliced_mode:
+            if self.type == "test" and self.sliced_mode:
                 instance = zoom(instance, (self.anisotropy_factor, 1, 1), order=0)
                 label = zoom(label, (self.anisotropy_factor, 1, 1), order=0)
-            instance = self.convert_zyx_to_czyx(instance, key='instance')  # CZYX
-            label = self.convert_zyx_to_czyx(label, key='label')  # CZYX
-            sample['instance'] = instance
-            sample['label'] = label
-        if (len(self.center_image_list) != 0):
-            if self.center_image_list[index][-3:] == 'csv':
-                center_image = self.rle_decode(self.center_image_list[index], center=True)
+            instance = self.convert_zyx_to_czyx(instance, key="instance")  # CZYX
+            label = self.convert_zyx_to_czyx(label, key="label")  # CZYX
+            sample["instance"] = instance
+            sample["label"] = label
+        if len(self.center_image_list) != 0:
+            if self.center_image_list[index][-3:] == "csv":
+                center_image = self.rle_decode(
+                    self.center_image_list[index], center=True
+                )
             else:
                 center_image = tifffile.imread(self.center_image_list[index])
-            if self.type == 'test' and self.sliced_mode:
-                center_image = zoom(center_image, (self.anisotropy_factor, 1, 1), order=0)
-            center_image = self.convert_zyx_to_czyx(center_image, key='center_image')  # CZYX
-            sample['center_image'] = center_image
+            if self.type == "test" and self.sliced_mode:
+                center_image = zoom(
+                    center_image, (self.anisotropy_factor, 1, 1), order=0
+                )
+            center_image = self.convert_zyx_to_czyx(
+                center_image, key="center_image"
+            )  # CZYX
+            sample["center_image"] = center_image
 
         # transform
-        if (self.transform is not None):
+        if self.transform is not None:
             return self.transform(sample)
         else:
             return sample
@@ -176,7 +236,9 @@ def __getitem__(self, index):
     @classmethod
     def decode_instance(cls, pic, one_hot, bg_id=None):
         pic = np.array(pic, copy=False, dtype=np.uint16)
-        instance_map = np.zeros((pic.shape[0], pic.shape[1], pic.shape[2]), dtype=np.int16)
+        instance_map = np.zeros(
+            (pic.shape[0], pic.shape[1], pic.shape[2]), dtype=np.int16
+        )
         class_map = np.zeros((pic.shape[0], pic.shape[1], pic.shape[2]), dtype=np.uint8)
 
         if bg_id is not None:
@@ -186,7 +248,7 @@ def decode_instance(cls, pic, one_hot, bg_id=None):
                 ids, _, _ = relabel_sequential(pic[mask])
 
                 instance_map[mask] = ids
-                if (one_hot):
+                if one_hot:
                     class_map[np.max(mask, axis=0)] = 1
                 else:
                     class_map[mask] = 1
@@ -202,20 +264,22 @@ def rle_decode(cls, filename, one_hot=False, center=False):
         if one_hot:
             mask_decoded = []
             for row in df_numpy:
-                d['counts'] = ast.literal_eval(row[1])
-                d['size'] = ast.literal_eval(row[2])
+                d["counts"] = ast.literal_eval(row[1])
+                d["size"] = ast.literal_eval(row[2])
                 mask = rletools.decode(d)  # returns binary mask
                 mask_decoded.append(mask)
         else:
             if center:
-                mask_decoded = np.zeros(ast.literal_eval(df_numpy[0][2]),
-                                        dtype=np.bool)  # obtain size by reading first row of csv file
+                mask_decoded = np.zeros(
+                    ast.literal_eval(df_numpy[0][2]), dtype=np.bool
+                )  # obtain size by reading first row of csv file
             else:
-                mask_decoded = np.zeros(ast.literal_eval(df_numpy[0][2]),
-                                        dtype=np.uint16)  # obtain size by reading first row of csv file
+                mask_decoded = np.zeros(
+                    ast.literal_eval(df_numpy[0][2]), dtype=np.uint16
+                )  # obtain size by reading first row of csv file
             for row in df_numpy:
-                d['counts'] = ast.literal_eval(row[1])
-                d['size'] = ast.literal_eval(row[2])
+                d["counts"] = ast.literal_eval(row[1])
+                d["size"] = ast.literal_eval(row[2])
                 mask = rletools.decode(d)  # returns binary mask
                 z, y, x = np.where(mask == 1)
                 mask_decoded[z, y, x] = int(row[0])
diff --git a/EmbedSeg/datasets/TwoDimensionalDataset.py b/EmbedSeg/datasets/TwoDimensionalDataset.py
index 6030a9e..d0ddbcc 100644
--- a/EmbedSeg/datasets/TwoDimensionalDataset.py
+++ b/EmbedSeg/datasets/TwoDimensionalDataset.py
@@ -9,79 +9,114 @@
 from skimage.segmentation import relabel_sequential
 from torch.utils.data import Dataset
 
-from EmbedSeg.utils.generate_crops import normalize_min_max_percentile, normalize_mean_std
+from EmbedSeg.utils.generate_crops import (
+    normalize_min_max_percentile,
+    normalize_mean_std,
+)
 
 
 class TwoDimensionalDataset(Dataset):
     """
-            A class used to create a PyTorch Dataset for handling 2D images and label masks
+    A class used to create a PyTorch Dataset for handling 2D images and label masks
 
-            Attributes
-            ----------
-            image_list: list of strings containing paths to the images
+    Attributes
+    ----------
+    image_list: list of strings containing paths to the images
 
-            instance_list: list of strings containing paths to the GT instances
+    instance_list: list of strings containing paths to the GT instances
 
-            center_image_list: list of strings containing paths to the center images
+    center_image_list: list of strings containing paths to the center images
 
-            bg_id: int
-                Id corresponding to the background, default=0
-            size: int
-                This is set equal to a different value than `real_size` in case the epoch is desired to be shorter
+    bg_id: int
+        Id corresponding to the background, default=0
+    size: int
+        This is set equal to a different value than `real_size`
+        in case the epoch is desired to be shorter
 
-            real_size: int
-                Actual number of the number of images
+    real_size: int
+        Actual number of the number of images
 
-            transform: PyTorch transform
+    transform: PyTorch transform
 
-            one_hot: bool
-                Should be set equal to True
+    one_hot: bool
+        Should be set equal to True
 
-            norm: str
-                Should be set equal to one of `min-max-percentile`, `absolute`, `mean-std`
+    norm: str
+        Should be set equal to one of `min-max-percentile`, `absolute`, `mean-std`
 
-            type: str
-                Should be set equal to one of `train`, `val` or `test`
-            data_type: str
-                Should be set equal to one of `8-bit` or `16-bit`
+    type: str
+        Should be set equal to one of `train`, `val` or `test`
+    data_type: str
+        Should be set equal to one of `8-bit` or `16-bit`
 
-            normalization: bool
-                Should be equal to True for test images and set equal to False during training (since crops are already normalized)
+    normalization: bool
+        Should be equal to True for test images and set equal to False
+        during training (since crops are already normalized)
 
-            Methods
-            -------
-            __init__: Initializes an object of class `TwoDimensionalDataset`
+    Methods
+    -------
+    __init__: Initializes an object of class `TwoDimensionalDataset`
 
-            __len__: Returns self.real_size if self.size = None
+    __len__: Returns self.real_size if self.size = None
 
-            convert_yx_to_cyx: Adds an additional dimension for channel
+    convert_yx_to_cyx: Adds an additional dimension for channel
 
-            decode_instance: In case, instances are available as tiffs, this method decodes them
+    decode_instance: In case, instances are available as tiffs,
+    this method decodes them
 
-            rle_decode: In case, instances are available as csv files (and not tiffs), this method decodes them
-            """
-
-    def __init__(self, data_dir='./', center='center-medoid', type="train", bg_id=0, size=None, transform=None,
-                 one_hot=False, norm='min-max-percentile', data_type='8-bit', normalization=False,
-                 anisotropy_factor=1.0, sliced_mode=False, uniform_ds_factor=1):
+    rle_decode: In case, instances are available as csv files (and not tiffs),
+    this method decodes them
+    """
 
-        print('2-D `{}` dataloader created! Accessing data from {}/{}/'.format(type, data_dir, type))
+    def __init__(
+        self,
+        data_dir="./",
+        center="center-medoid",
+        type="train",
+        bg_id=0,
+        size=None,
+        transform=None,
+        one_hot=False,
+        norm="min-max-percentile",
+        data_type="8-bit",
+        normalization=False,
+        anisotropy_factor=1.0,
+        sliced_mode=False,
+        uniform_ds_factor=1,
+    ):
+        print(
+            "2-D `{}` dataloader created! Accessing data from {}/{}/".format(
+                type, data_dir, type
+            )
+        )
 
         # get image and instance list
-        image_list = glob.glob(os.path.join(data_dir, '{}/'.format(type), 'images/*.tif'))
+        image_list = glob.glob(
+            os.path.join(data_dir, "{}/".format(type), "images/*.tif")
+        )
         image_list.sort()
-        print('Number of images in `{}` directory is {}'.format(type, len(image_list)))
+        print("Number of images in `{}` directory is {}".format(type, len(image_list)))
         self.image_list = image_list
 
-        instance_list = glob.glob(os.path.join(data_dir, '{}/'.format(type), 'masks/*'))
+        instance_list = glob.glob(os.path.join(data_dir, "{}/".format(type), "masks/*"))
         instance_list.sort()
-        print('Number of instances in `{}` directory is {}'.format(type, len(instance_list)))
+        print(
+            "Number of instances in `{}` directory is {}".format(
+                type, len(instance_list)
+            )
+        )
         self.instance_list = instance_list
 
-        center_image_list = glob.glob(os.path.join(data_dir, '{}/'.format(type), center + '/*'))
+        center_image_list = glob.glob(
+            os.path.join(data_dir, "{}/".format(type), center + "/*")
+        )
         center_image_list.sort()
-        print('Number of center images in `{}` directory is {}'.format(type, len(center_image_list)))
-        print('*************************')
+        print(
+            "Number of center images in `{}` directory is {}".format(
+                type, len(center_image_list)
+            )
+        )
+        print("*************************")
         self.center_image_list = center_image_list
 
         self.bg_id = bg_id
@@ -95,65 +130,71 @@ def __init__(self, data_dir='./', center='center-medoid', type="train", bg_id=0,
         self.normalization = normalization
 
     def __len__(self):
-
         return self.real_size if self.size is None else self.size
 
     def convert_yx_to_cyx(self, im, key):
-        if im.ndim == 2 and key == 'image':  # gray-scale image
+        if im.ndim == 2 and key == "image":  # gray-scale image
             im = im[np.newaxis, ...]  # CYX
-        elif im.ndim == 3 and key == 'image':  # multi-channel image image
+        elif im.ndim == 3 and key == "image":  # multi-channel image image
             pass
         else:
             im = im[np.newaxis, ...]
         return im
 
     def __getitem__(self, index):
-
         index = index if self.size is None else random.randint(0, self.real_size - 1)
         sample = {}
 
         # load image
         image = tifffile.imread(self.image_list[index])  # YX or CYX
-        if (self.normalization and self.norm == 'min-max-percentile'):
+        if self.normalization and self.norm == "min-max-percentile":
             if image.ndim == 2:  # gray-scale
                 image = normalize_min_max_percentile(image, 1, 99.8, axis=(0, 1))
             elif image.ndim == 3:  # multi-channel image (C, H, W)
                 image = normalize_min_max_percentile(image, 1, 99.8, axis=(1, 2))
-        elif (self.normalization and self.norm == 'mean-std'):
+        elif self.normalization and self.norm == "mean-std":
             if image.ndim == 2:
                 image = normalize_mean_std(image)  # axis == None
             elif image.ndim == 3:
                 image = normalize_mean_std(image, axis=(1, 2))
-        elif (self.normalization and self.norm == 'absolute'):
-            image = image.astype(np.float32)
-            if self.data_type == '8-bit':
+        elif self.normalization and self.norm == "absolute":
+            image = image.astype(float)
+            if self.data_type == "8-bit":
                 image /= 255
-            elif self.data_type == '16-bit':
+            elif self.data_type == "16-bit":
                 image /= 65535
-        image = self.convert_yx_to_cyx(image, key='image')
-        sample['image'] = image  # CYX
-        sample['im_name'] = self.image_list[index]
-
-        if (len(self.instance_list) != 0):
-            if self.instance_list[index][-3:] == 'csv':
-                instance = self.rle_decode(self.instance_list[index], one_hot=self.one_hot)
+        image = self.convert_yx_to_cyx(image, key="image")
+        sample["image"] = image  # CYX
+        sample["im_name"] = self.image_list[index]
+
+        if len(self.instance_list) != 0:
+            if self.instance_list[index][-3:] == "csv":
+                instance = self.rle_decode(
+                    self.instance_list[index], one_hot=self.one_hot
+                )
             else:
-                instance = tifffile.imread(self.instance_list[index])  # YX or DYX (one-hot!)
+                instance = tifffile.imread(
+                    self.instance_list[index]
+                )  # YX or DYX (one-hot!)
             instance, label = self.decode_instance(instance, self.one_hot, self.bg_id)
-            instance = self.convert_yx_to_cyx(instance, key='instance')  # CYX or CDYX
-            label = self.convert_yx_to_cyx(label, key='label')  # CYX
-            sample['instance'] = instance
-            sample['label'] = label
-        if (len(self.center_image_list) != 0):
-            if self.center_image_list[index][-3:] == 'csv':
-                center_image = self.rle_decode(self.center_image_list[index], center=True)
+            instance = self.convert_yx_to_cyx(instance, key="instance")  # CYX or CDYX
+            label = self.convert_yx_to_cyx(label, key="label")  # CYX
+            sample["instance"] = instance
+            sample["label"] = label
+        if len(self.center_image_list) != 0:
+            if self.center_image_list[index][-3:] == "csv":
+                center_image = self.rle_decode(
+                    self.center_image_list[index], center=True
+                )
             else:
                 center_image = tifffile.imread(self.center_image_list[index])
-            center_image = self.convert_yx_to_cyx(center_image, key='center_image')  # CYX
-            sample['center_image'] = center_image
+            center_image = self.convert_yx_to_cyx(
+                center_image, key="center_image"
+            )  # CYX
+            sample["center_image"] = center_image
 
         # transform
-        if (self.transform is not None):
+        if self.transform is not None:
             return self.transform(sample)
         else:
             return sample
@@ -161,8 +202,10 @@ def __getitem__(self, index):
     @classmethod
     def decode_instance(cls, pic, one_hot, bg_id=None):
         pic = np.array(pic, copy=False, dtype=np.uint16)
-        if (one_hot):
-            instance_map = np.zeros((pic.shape[0], pic.shape[1], pic.shape[2]), dtype=np.uint8)
+        if one_hot:
+            instance_map = np.zeros(
+                (pic.shape[0], pic.shape[1], pic.shape[2]), dtype=np.uint8
+            )
             class_map = np.zeros((pic.shape[1], pic.shape[2]), dtype=np.uint8)
         else:
             instance_map = np.zeros((pic.shape[0], pic.shape[1]), dtype=np.int16)
@@ -173,7 +216,7 @@ def decode_instance(cls, pic, one_hot, bg_id=None):
             if mask.sum() > 0:
                 ids, _, _ = relabel_sequential(pic[mask])
                 instance_map[mask] = ids
-                if (one_hot):
+                if one_hot:
                     class_map[np.max(mask, axis=0)] = 1
                 else:
                     class_map[mask] = 1
@@ -189,20 +232,22 @@ def rle_decode(cls, filename, one_hot=False, center=False):
         if one_hot:
             mask_decoded = []
             for row in df_numpy:
-                d['counts'] = ast.literal_eval(row[1])
-                d['size'] = ast.literal_eval(row[2])
+                d["counts"] = ast.literal_eval(row[1])
+                d["size"] = ast.literal_eval(row[2])
                 mask = rletools.decode(d)  # returns binary mask
                 mask_decoded.append(mask)
         else:
             if center:
-                mask_decoded = np.zeros(ast.literal_eval(df_numpy[0][2]),
-                                        dtype=np.bool)  # obtain size by reading first row of csv file
+                mask_decoded = np.zeros(
+                    ast.literal_eval(df_numpy[0][2]), dtype=np.bool
+                )  # obtain size by reading first row of csv file
             else:
-                mask_decoded = np.zeros(ast.literal_eval(df_numpy[0][2]),
-                                        dtype=np.uint16)  # obtain size by reading first row of csv file
+                mask_decoded = np.zeros(
+                    ast.literal_eval(df_numpy[0][2]), dtype=np.uint16
+                )  # obtain size by reading first row of csv file
             for row in df_numpy:
-                d['counts'] = ast.literal_eval(row[1])
-                d['size'] = ast.literal_eval(row[2])
+                d["counts"] = ast.literal_eval(row[1])
+                d["size"] = ast.literal_eval(row[2])
                 mask = rletools.decode(d)  # returns binary mask
                 y, x = np.where(mask == 1)
                 mask_decoded[y, x] = int(row[0])
diff --git a/EmbedSeg/datasets/__init__.py b/EmbedSeg/datasets/__init__.py
index 9d29d12..afd7ee6 100644
--- a/EmbedSeg/datasets/__init__.py
+++ b/EmbedSeg/datasets/__init__.py
@@ -1,14 +1,15 @@
 from EmbedSeg.datasets.TwoDimensionalDataset import TwoDimensionalDataset
 from EmbedSeg.datasets.ThreeDimensionalDataset import ThreeDimensionalDataset
 
+
 def get_dataset(name, dataset_opts):
-    if name == "2d": 
+    if name == "2d":
         return TwoDimensionalDataset(**dataset_opts)
-    elif name =="3d_sliced":
+    elif name == "3d_sliced":
         return ThreeDimensionalDataset(**dataset_opts)
-    elif name=="3d":
+    elif name == "3d":
         return ThreeDimensionalDataset(**dataset_opts)
-    elif name =="3d_ilp":
+    elif name == "3d_ilp":
         return ThreeDimensionalDataset(**dataset_opts)
     else:
         raise RuntimeError("Dataset {} not available".format(name))
diff --git a/EmbedSeg/models/BranchedERFNet.py b/EmbedSeg/models/BranchedERFNet.py
index dd1e423..a7b59ea 100644
--- a/EmbedSeg/models/BranchedERFNet.py
+++ b/EmbedSeg/models/BranchedERFNet.py
@@ -1,6 +1,7 @@
 """
 Author: Davy Neven
-Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
+Licensed under the CC BY-NC 4.0 license 
+(https://creativecommons.org/licenses/by-nc/4.0/)
 """
 import torch
 import torch.nn as nn
@@ -8,17 +9,16 @@
 
 
 class BranchedERFNet(nn.Module):
-    def __init__(self, num_classes, input_channels = 1, encoder=None):
+    def __init__(self, num_classes, input_channels=1, encoder=None):
         super().__init__()
 
-        print('Creating Branched Erfnet with {} outputs'.format(num_classes))
-        
-        if (encoder is None):
+        print("Creating Branched Erfnet with {} outputs".format(num_classes))
+
+        if encoder is None:
             self.encoder = erfnet.Encoder(sum(num_classes), input_channels)
-            
+
         else:
             self.encoder = encoder
-            
 
         self.decoders = nn.ModuleList()
         for n in num_classes:
@@ -27,14 +27,13 @@ def __init__(self, num_classes, input_channels = 1, encoder=None):
     def init_output(self, n_sigma=1):
         with torch.no_grad():
             output_conv = self.decoders[0].output_conv
-            print('Initialize last layer with size: ',
-                  output_conv.weight.size())
-            print('*************************')
+            print("Initialize last layer with size: ", output_conv.weight.size())
+            print("*************************")
             output_conv.weight[:, 0:2, :, :].fill_(0)
             output_conv.bias[0:2].fill_(0)
 
-            output_conv.weight[:, 2:2+n_sigma, :, :].fill_(0)
-            output_conv.bias[2:2+n_sigma].fill_(1)
+            output_conv.weight[:, 2 : 2 + n_sigma, :, :].fill_(0)
+            output_conv.bias[2 : 2 + n_sigma].fill_(1)
 
     def forward(self, input, only_encode=False):
         if only_encode:
diff --git a/EmbedSeg/models/BranchedERFNet_3d.py b/EmbedSeg/models/BranchedERFNet_3d.py
index 6a7e2ad..2fde04b 100644
--- a/EmbedSeg/models/BranchedERFNet_3d.py
+++ b/EmbedSeg/models/BranchedERFNet_3d.py
@@ -4,11 +4,11 @@
 
 
 class BranchedERFNet_3d(nn.Module):
-    def __init__(self, num_classes, input_channels =1, encoder=None):
+    def __init__(self, num_classes, input_channels=1, encoder=None):
         super().__init__()
 
-        print('Creating Branched Erfnet 3D with {} outputs'.format(num_classes))
-        if (encoder is None):
+        print("Creating Branched Erfnet 3D with {} outputs".format(num_classes))
+        if encoder is None:
             self.encoder = erfnet.Encoder(sum(num_classes), input_channels)
         else:
             self.encoder = encoder
@@ -20,14 +20,13 @@ def __init__(self, num_classes, input_channels =1, encoder=None):
     def init_output(self, n_sigma=1):
         with torch.no_grad():
             output_conv = self.decoders[0].output_conv
-            print('initialize last layer with size: ',
-                  output_conv.weight.size())
+            print("initialize last layer with size: ", output_conv.weight.size())
 
             output_conv.weight[:, 0:3, :, :, :].fill_(0)
             output_conv.bias[0:3].fill_(0)
 
-            output_conv.weight[:, 3:3 + n_sigma, :, :, :].fill_(0)
-            output_conv.bias[3:3 + n_sigma].fill_(1)
+            output_conv.weight[:, 3 : 3 + n_sigma, :, :, :].fill_(0)
+            output_conv.bias[3 : 3 + n_sigma].fill_(1)
 
     def forward(self, input, only_encode=False):
         if only_encode:
@@ -36,5 +35,3 @@ def forward(self, input, only_encode=False):
             output = self.encoder(input)
 
         return torch.cat([decoder.forward(output) for decoder in self.decoders], 1)
-
-
diff --git a/EmbedSeg/models/__init__.py b/EmbedSeg/models/__init__.py
index 9e2bf5b..b3df841 100644
--- a/EmbedSeg/models/__init__.py
+++ b/EmbedSeg/models/__init__.py
@@ -1,11 +1,13 @@
 from EmbedSeg.models.BranchedERFNet import BranchedERFNet
 from EmbedSeg.models.BranchedERFNet_3d import BranchedERFNet_3d
+
+
 def get_model(name, model_opts):
     if name == "branched_erfnet":
         model = BranchedERFNet(**model_opts)
         return model
-    elif name=="branched_erfnet_3d":
+    elif name == "branched_erfnet_3d":
         model = BranchedERFNet_3d(**model_opts)
         return model
     else:
-        raise RuntimeError("model \"{}\" not available".format(name))
\ No newline at end of file
+        raise RuntimeError('model "{}" not available'.format(name))
diff --git a/EmbedSeg/models/erfnet.py b/EmbedSeg/models/erfnet.py
index d58c18b..ccb91ec 100644
--- a/EmbedSeg/models/erfnet.py
+++ b/EmbedSeg/models/erfnet.py
@@ -5,16 +5,16 @@
 
 import torch
 import torch.nn as nn
-import torch.nn.init as init
 import torch.nn.functional as F
 
 
-class DownsamplerBlock (nn.Module):
+class DownsamplerBlock(nn.Module):
     def __init__(self, ninput, noutput):
         super().__init__()
 
-        self.conv = nn.Conv2d(ninput, noutput-ninput,
-                              (3, 3), stride=2, padding=1, bias=True)
+        self.conv = nn.Conv2d(
+            ninput, noutput - ninput, (3, 3), stride=2, padding=1, bias=True
+        )
         self.pool = nn.MaxPool2d(2, stride=2)
         self.bn = nn.BatchNorm2d(noutput, eps=1e-3)
 
@@ -24,30 +24,45 @@ def forward(self, input):
         return F.relu(output)
 
 
-class non_bottleneck_1d (nn.Module):
+class non_bottleneck_1d(nn.Module):
     def __init__(self, chann, dropprob, dilated):
         super().__init__()
 
         self.conv3x1_1 = nn.Conv2d(
-            chann, chann, (3, 1), stride=1, padding=(1, 0), bias=True)
+            chann, chann, (3, 1), stride=1, padding=(1, 0), bias=True
+        )
 
         self.conv1x3_1 = nn.Conv2d(
-            chann, chann, (1, 3), stride=1, padding=(0, 1), bias=True)
+            chann, chann, (1, 3), stride=1, padding=(0, 1), bias=True
+        )
 
         self.bn1 = nn.BatchNorm2d(chann, eps=1e-03)
 
-        self.conv3x1_2 = nn.Conv2d(chann, chann, (3, 1), stride=1, padding=(
-            1*dilated, 0), bias=True, dilation=(dilated, 1))
-
-        self.conv1x3_2 = nn.Conv2d(chann, chann, (1, 3), stride=1, padding=(
-            0, 1*dilated), bias=True, dilation=(1, dilated))
+        self.conv3x1_2 = nn.Conv2d(
+            chann,
+            chann,
+            (3, 1),
+            stride=1,
+            padding=(1 * dilated, 0),
+            bias=True,
+            dilation=(dilated, 1),
+        )
+
+        self.conv1x3_2 = nn.Conv2d(
+            chann,
+            chann,
+            (1, 3),
+            stride=1,
+            padding=(0, 1 * dilated),
+            bias=True,
+            dilation=(1, dilated),
+        )
 
         self.bn2 = nn.BatchNorm2d(chann, eps=1e-03)
 
         self.dropout = nn.Dropout2d(dropprob)
 
     def forward(self, input):
-
         output = self.conv3x1_1(input)
         output = F.relu(output)
         output = self.conv1x3_1(output)
@@ -59,16 +74,18 @@ def forward(self, input):
         output = self.conv1x3_2(output)
         output = self.bn2(output)
 
-        if (self.dropout.p != 0):
+        if self.dropout.p != 0:
             output = self.dropout(output)
 
-        return F.relu(output+input)  # +input = identity (residual connection)
+        return F.relu(output + input)  # +input = identity (residual connection)
 
 
 class Encoder(nn.Module):
     def __init__(self, num_classes, input_channels):
         super().__init__()
-        self.initial_block = DownsamplerBlock(input_channels, 16) # TODO input_channels = 1 (for gray-scale), 3 (for RGB)
+        self.initial_block = DownsamplerBlock(
+            input_channels, 16
+        )  # TODO input_channels = 1 (for gray-scale), 3 (for RGB)
         self.layers = nn.ModuleList()
 
         self.layers.append(DownsamplerBlock(16, 64))
@@ -86,7 +103,8 @@ def __init__(self, num_classes, input_channels):
 
         # Only in encoder mode:
         self.output_conv = nn.Conv2d(
-            128, num_classes, 1, stride=1, padding=0, bias=True)
+            128, num_classes, 1, stride=1, padding=0, bias=True
+        )
 
     def forward(self, input, predict=False):
         output = self.initial_block(input)
@@ -100,11 +118,12 @@ def forward(self, input, predict=False):
         return output
 
 
-class UpsamplerBlock (nn.Module):
+class UpsamplerBlock(nn.Module):
     def __init__(self, ninput, noutput):
         super().__init__()
         self.conv = nn.ConvTranspose2d(
-            ninput, noutput, 3, stride=2, padding=1, output_padding=1, bias=True)
+            ninput, noutput, 3, stride=2, padding=1, output_padding=1, bias=True
+        )
         self.bn = nn.BatchNorm2d(noutput, eps=1e-3)
 
     def forward(self, input):
@@ -113,7 +132,7 @@ def forward(self, input):
         return F.relu(output)
 
 
-class Decoder (nn.Module):
+class Decoder(nn.Module):
     def __init__(self, num_classes):
         super().__init__()
 
@@ -128,7 +147,8 @@ def __init__(self, num_classes):
         self.layers.append(non_bottleneck_1d(16, 0, 1))
 
         self.output_conv = nn.ConvTranspose2d(
-            16, num_classes, 2, stride=2, padding=0, output_padding=0, bias=True)
+            16, num_classes, 2, stride=2, padding=0, output_padding=0, bias=True
+        )
 
     def forward(self, input):
         output = input
@@ -140,14 +160,17 @@ def forward(self, input):
 
         return output
 
+
 # ERFNet
 
 
 class Net(nn.Module):
-    def __init__(self, num_classes, input_channels, encoder=None):  # use encoder to pass pretrained encoder
+    def __init__(
+        self, num_classes, input_channels, encoder=None
+    ):  # use encoder to pass pretrained encoder
         super().__init__()
 
-        if (encoder == None):
+        if encoder is None:
             self.encoder = Encoder(num_classes, input_channels)
         else:
             self.encoder = encoder
diff --git a/EmbedSeg/models/erfnet_3d.py b/EmbedSeg/models/erfnet_3d.py
index e9e0c69..42c267f 100644
--- a/EmbedSeg/models/erfnet_3d.py
+++ b/EmbedSeg/models/erfnet_3d.py
@@ -7,12 +7,13 @@ class DownsamplerBlock(nn.Module):
     def __init__(self, ninput, noutput):
         super().__init__()
 
-        self.conv = nn.Conv3d(ninput, noutput - ninput,
-                              (3, 3, 3), stride=2, padding=1, bias=True)  
+        self.conv = nn.Conv3d(
+            ninput, noutput - ninput, (3, 3, 3), stride=2, padding=1, bias=True
+        )
 
-        self.pool = nn.MaxPool3d(2, stride=2)  
+        self.pool = nn.MaxPool3d(2, stride=2)
 
-        self.bn = nn.BatchNorm3d(noutput, eps=1e-3)  
+        self.bn = nn.BatchNorm3d(noutput, eps=1e-3)
 
     def forward(self, input):
         output = torch.cat([self.conv(input), self.pool(input)], 1)
@@ -24,24 +25,48 @@ class non_bottleneck_1d(nn.Module):
     def __init__(self, chann, dropprob, dilated):
         super().__init__()
         self.conv3x1x1_1 = nn.Conv3d(
-            chann, chann, (3, 1, 1), stride=1, padding=(1, 0, 0), bias=True)  
+            chann, chann, (3, 1, 1), stride=1, padding=(1, 0, 0), bias=True
+        )
 
         self.conv1x3x1_1 = nn.Conv3d(
-            chann, chann, (1, 3, 1), stride=1, padding=(0, 1, 0), bias=True)  
+            chann, chann, (1, 3, 1), stride=1, padding=(0, 1, 0), bias=True
+        )
 
         self.conv1x1x3_1 = nn.Conv3d(
-            chann, chann, (1, 1, 3), stride=1, padding=(0, 0, 1), bias=True)  
-
-        self.bn1 = nn.BatchNorm3d(chann, eps=1e-03)  
-
-        self.conv3x1x1_2 = nn.Conv3d(chann, chann, (3, 1, 1), stride=1, padding=(
-            1 * dilated, 0, 0), bias=True, dilation=(dilated, 1, 1))  
-
-        self.conv1x3x1_2 = nn.Conv3d(chann, chann, (1, 3, 1), stride=1, padding=(
-            0, 1 * dilated, 0), bias=True, dilation=(1, dilated, 1))  
-
-        self.conv1x1x3_2 = nn.Conv3d(chann, chann, (1, 1, 3), stride=1, padding=(
-            0, 0, 1 * dilated), bias=True, dilation=(1, 1, dilated))
+            chann, chann, (1, 1, 3), stride=1, padding=(0, 0, 1), bias=True
+        )
+
+        self.bn1 = nn.BatchNorm3d(chann, eps=1e-03)
+
+        self.conv3x1x1_2 = nn.Conv3d(
+            chann,
+            chann,
+            (3, 1, 1),
+            stride=1,
+            padding=(1 * dilated, 0, 0),
+            bias=True,
+            dilation=(dilated, 1, 1),
+        )
+
+        self.conv1x3x1_2 = nn.Conv3d(
+            chann,
+            chann,
+            (1, 3, 1),
+            stride=1,
+            padding=(0, 1 * dilated, 0),
+            bias=True,
+            dilation=(1, dilated, 1),
+        )
+
+        self.conv1x1x3_2 = nn.Conv3d(
+            chann,
+            chann,
+            (1, 1, 3),
+            stride=1,
+            padding=(0, 0, 1 * dilated),
+            bias=True,
+            dilation=(1, 1, dilated),
+        )
 
         self.bn2 = nn.BatchNorm3d(chann, eps=1e-03)
 
@@ -50,7 +75,7 @@ def __init__(self, chann, dropprob, dilated):
     def forward(self, input):
         output = self.conv3x1x1_1(input)
         output = F.relu(output)
-        output = self.conv1x3x1_1(output)  
+        output = self.conv1x3x1_1(output)
         output = F.relu(output)
         output = self.conv1x1x3_1(output)
         # TODO think about adding a relu here!
@@ -60,12 +85,12 @@ def forward(self, input):
 
         output = self.conv3x1x1_2(output)
         output = F.relu(output)
-        output = self.conv1x3x1_2(output)  
+        output = self.conv1x3x1_2(output)
         output = F.relu(output)
         output = self.conv1x1x3_2(output)
         output = self.bn2(output)
 
-        if (self.dropout.p != 0):
+        if self.dropout.p != 0:
             output = self.dropout(output)
 
         return F.relu(output + input)  # +input = identity (residual connection)
@@ -74,7 +99,7 @@ def forward(self, input):
 class Encoder(nn.Module):
     def __init__(self, num_classes, input_channels):
         super().__init__()
-        self.initial_block = DownsamplerBlock(input_channels, 16)  
+        self.initial_block = DownsamplerBlock(input_channels, 16)
 
         self.layers = nn.ModuleList()
 
@@ -92,7 +117,8 @@ def __init__(self, num_classes, input_channels):
             self.layers.append(non_bottleneck_1d(128, 0.3, 16))
 
         self.output_conv = nn.Conv3d(
-            128, num_classes, 1, stride=1, padding=0, bias=True)  
+            128, num_classes, 1, stride=1, padding=0, bias=True
+        )
 
     def forward(self, input, predict=False):
         output = self.initial_block(input)
@@ -110,8 +136,9 @@ class UpsamplerBlock(nn.Module):
     def __init__(self, ninput, noutput):
         super().__init__()
         self.conv = nn.ConvTranspose3d(
-            ninput, noutput, 3, stride=2, padding=1, output_padding=1, bias=True)
-        self.bn = nn.BatchNorm3d(noutput, eps=1e-3)  
+            ninput, noutput, 3, stride=2, padding=1, output_padding=1, bias=True
+        )
+        self.bn = nn.BatchNorm3d(noutput, eps=1e-3)
 
     def forward(self, input):
         output = self.conv(input)
@@ -134,7 +161,8 @@ def __init__(self, num_classes):
         self.layers.append(non_bottleneck_1d(16, 0, 1))
 
         self.output_conv = nn.ConvTranspose3d(
-            16, num_classes, 2, stride=2, padding=0, output_padding=0, bias=True)
+            16, num_classes, 2, stride=2, padding=0, output_padding=0, bias=True
+        )
 
     def forward(self, input):
         output = input
@@ -149,10 +177,12 @@ def forward(self, input):
 
 # ERFNet
 class Net(nn.Module):
-    def __init__(self, num_classes, input_channels, encoder=None):  # use encoder to pass pretrained encoder
+    def __init__(
+        self, num_classes, input_channels, encoder=None
+    ):  # use encoder to pass pretrained encoder
         super().__init__()
 
-        if (encoder == None):
+        if encoder is None:
             self.encoder = Encoder(num_classes, input_channels)
         else:
             self.encoder = encoder
diff --git a/EmbedSeg/train.py b/EmbedSeg/train.py
index 67819dd..0954c99 100644
--- a/EmbedSeg/train.py
+++ b/EmbedSeg/train.py
@@ -8,25 +8,35 @@
 from EmbedSeg.criterions import get_loss
 from EmbedSeg.datasets import get_dataset
 from EmbedSeg.models import get_model
-from EmbedSeg.utils.utils import AverageMeter, Cluster, Cluster_3d, Logger, Visualizer, \
-    prepare_embedding_for_train_image
+from EmbedSeg.utils.utils import (
+    AverageMeter,
+    Cluster,
+    Cluster_3d,
+    Logger,
+    Visualizer,
+    prepare_embedding_for_train_image,
+)
 
-torch.backends.cudnn.benchmark = True
 import numpy as np
 
+torch.backends.cudnn.benchmark = True
+
 
-def train(virtual_batch_multiplier, one_hot, n_sigma, args):
+def train(virtual_batch_multiplier, one_hot, n_sigma, args, device):
     """Trains 2D Model with virtual multiplier
 
     Virtual batching code inspired from:
-    https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
+    https://medium.com/huggingface/training-larger-batches-
+    practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
 
     Parameters
     ----------
     virtual_batch_multiplier : int
-        Set to 1 by default. Effective batch size is product of virtual_batch_multiplier and batch-size
+        Set to 1 by default. Effective batch size is product of
+        virtual_batch_multiplier and batch-size
     one_hot : bool
-        In case the GT labels are available in one-hot fashion, this parameter is set equal to True
+        In case the GT labels are available in one-hot fashion,
+        this parameter is set equal to True
     n_sigma: int
         Should be equal to 2 for a 2D model
     args: dictionary
@@ -42,14 +52,15 @@ def train(virtual_batch_multiplier, one_hot, n_sigma, args):
     # put model into training mode
     model.train()
     for param_group in optimizer.param_groups:
-        print('learning rate: {}'.format(param_group['lr']))
+        print("learning rate: {}".format(param_group["lr"]))
 
     optimizer.zero_grad()  # Reset gradients tensors
     for i, sample in enumerate(tqdm(train_dataset_it)):
-        im = sample['image']
-        instances = sample['instance'].squeeze(1)
-        class_labels = sample['label'].squeeze(1)
-        center_images = sample['center_image'].squeeze(1)
+        im = sample["image"]
+        im = im.to(device)
+        instances = sample["instance"].squeeze(1).to(device)
+        class_labels = sample["label"].squeeze(1).to(device)
+        center_images = sample["center_image"].squeeze(1).to(device)
         output = model(im)  # Forward pass
         loss = criterion(output, instances, class_labels, center_images, **args)
         loss = loss / virtual_batch_multiplier  # Normalize our loss (if averaged)
@@ -62,54 +73,71 @@ def train(virtual_batch_multiplier, one_hot, n_sigma, args):
     return loss_meter.avg * virtual_batch_multiplier
 
 
-def train_vanilla(display, display_embedding, display_it, one_hot, grid_x, grid_y, pixel_x, pixel_y, n_sigma,
-                  args):
+def train_vanilla(
+    display,
+    display_embedding,
+    display_it,
+    one_hot,
+    grid_x,
+    grid_y,
+    pixel_x,
+    pixel_y,
+    n_sigma,
+    args,
+    device,
+):
     """Trains 2D Model without virtual multiplier.
 
-        Parameters
-        ----------
-        display : bool
-            Displays input, GT, model predictions during training
-        display_embedding : bool
-            Displays embeddings for train (crop) images
-        display_it: int
-                Displays a new training image, the corresponding GT and model prediction every `display_it` crop images
-        one_hot: bool
-            In case the GT labels are available in one-hot fashion, this parameter is set equal to True
-        grid_x: int
-            Number of pixels along x dimension which constitute a tile
-        grid_y: int
-            Number of pixels along y dimension which constitute a tile
-        pixel_x: float
-            The grid length along x is mapped to `pixel_x`.
-            For example, if grid_x = 1024 and pixel_x = 1.0, then each dX or the spacing between consecutive pixels
-            along the x dimension is set equal to pixel_x/grid_x = 1.0/1024
-        pixel_y: float
-            The grid length along y is mapped to `pixel_y`.
-            For example, if grid_y = 1024 and pixel_y = 1.0, then each dY or the spacing between consecutive pixels
-            along the y dimension is set equal to pixel_y/grid_y = 1.0/1024
-        n_sigma: int
-            Should be set equal to 2 for a 2D model
-        args: dictionary
-
-        Returns
-        -------
-        float
-            Average loss
-        """
+    Parameters
+    ----------
+    display : bool
+        Displays input, GT, model predictions during training
+    display_embedding : bool
+        Displays embeddings for train (crop) images
+    display_it: int
+        Displays a new training image, the corresponding GT and model
+        prediction every `display_it` crop images
+    one_hot: bool
+        In case the GT labels are available in one-hot fashion,
+        this parameter is set equal to True
+    grid_x: int
+        Number of pixels along x dimension which constitute a tile
+    grid_y: int
+        Number of pixels along y dimension which constitute a tile
+    pixel_x: float
+        The grid length along x is mapped to `pixel_x`.
+        For example, if grid_x = 1024 and pixel_x = 1.0,
+        then each dX or the spacing between consecutive pixels
+        along the x dimension is set equal to pixel_x/grid_x = 1.0/1024
+    pixel_y: float
+        The grid length along y is mapped to `pixel_y`.
+        For example, if grid_y = 1024 and pixel_y = 1.0,
+        then each dY or the spacing between consecutive pixels
+        along the y dimension is set equal to pixel_y/grid_y = 1.0/1024
+    n_sigma: int
+        Should be set equal to 2 for a 2D model
+    args: dictionary
+
+    Returns
+    -------
+    float
+        Average loss
+    """
     # define meters
     loss_meter = AverageMeter()
     # put model into training mode
     model.train()
 
     for param_group in optimizer.param_groups:
-        print('learning rate: {}'.format(param_group['lr']))
+        print("learning rate: {}".format(param_group["lr"]))
     for i, sample in enumerate(tqdm(train_dataset_it)):
-
-        im = sample['image']
-        instances = sample['instance'].squeeze(1)  # 1YX (not one-hot) or 1DYX (one-hot)
-        class_labels = sample['label'].squeeze(1)  # 1YX
-        center_images = sample['center_image'].squeeze(1)  # 1YX
+        im = sample["image"]
+        im = im.to(device)
+        instances = (
+            sample["instance"].squeeze(1).to(device)
+        )  # 1YX (not one-hot) or 1DYX (one-hot)
+        class_labels = sample["label"].squeeze(1).to(device)  # 1YX
+        center_images = sample["center_image"].squeeze(1).to(device)  # 1YX
         output = model(im)  # B 5 Y X
         loss = criterion(output, instances, class_labels, center_images, **args)
         loss = loss.mean()
@@ -119,154 +147,225 @@ def train_vanilla(display, display_embedding, display_it, one_hot, grid_x, grid_
         loss_meter.update(loss.item())
         if display and i % display_it == 0:
             with torch.no_grad():
-                visualizer.display(im[0], key='image', title='Image')
-                predictions = cluster.cluster_with_gt(output[0], instances[0], n_sigma=n_sigma)
+                visualizer.display(im[0], key="image", title="Image")
+                predictions = cluster.cluster_with_gt(
+                    output[0], instances[0], n_sigma=n_sigma
+                )
                 if one_hot:
                     instance = invert_one_hot(instances[0].cpu().detach().numpy())
-                    visualizer.display(instance, key='groundtruth', title='Ground Truth')  # TODO
+                    visualizer.display(
+                        instance, key="groundtruth", title="Ground Truth"
+                    )  # TODO
                     instance_ids = np.arange(instances.size(1))  # instances[0] --> DYX
                 else:
-                    visualizer.display(instances[0].cpu(), key='groundtruth', title='Ground Truth')  # TODO
+                    visualizer.display(
+                        instances[0].cpu(), key="groundtruth", title="Ground Truth"
+                    )  # TODO
                     instance_ids = instances[0].unique()
                     instance_ids = instance_ids[instance_ids != 0]
 
                 if display_embedding:
-                    center_x, center_y, samples_x, samples_y, sample_spatial_embedding_x, \
-                    sample_spatial_embedding_y, sigma_x, sigma_y, color_sample_dic, color_embedding_dic = \
-                        prepare_embedding_for_train_image(one_hot=one_hot, grid_x=grid_x, grid_y=grid_y,
-                                                          pixel_x=pixel_x, pixel_y=pixel_y,
-                                                          predictions=predictions, instance_ids=instance_ids,
-                                                          center_images=center_images,
-                                                          output=output, instances=instances, n_sigma=n_sigma)
+                    (
+                        center_x,
+                        center_y,
+                        samples_x,
+                        samples_y,
+                        sample_spatial_embedding_x,
+                        sample_spatial_embedding_y,
+                        sigma_x,
+                        sigma_y,
+                        color_sample_dic,
+                        color_embedding_dic,
+                    ) = prepare_embedding_for_train_image(
+                        one_hot=one_hot,
+                        grid_x=grid_x,
+                        grid_y=grid_y,
+                        pixel_x=pixel_x,
+                        pixel_y=pixel_y,
+                        predictions=predictions,
+                        instance_ids=instance_ids,
+                        center_images=center_images,
+                        output=output,
+                        instances=instances,
+                        n_sigma=n_sigma,
+                    )
                     if one_hot:
-                        visualizer.display(torch.max(instances[0], dim=0)[0], key='center', title='Center',
-                                           center_x=center_x,
-                                           center_y=center_y,
-                                           samples_x=samples_x, samples_y=samples_y,
-                                           sample_spatial_embedding_x=sample_spatial_embedding_x,
-                                           sample_spatial_embedding_y=sample_spatial_embedding_y,
-                                           sigma_x=sigma_x, sigma_y=sigma_y,
-                                           color_sample=color_sample_dic, color_embedding=color_embedding_dic)
+                        visualizer.display(
+                            torch.max(instances[0], dim=0)[0],
+                            key="center",
+                            title="Center",
+                            center_x=center_x,
+                            center_y=center_y,
+                            samples_x=samples_x,
+                            samples_y=samples_y,
+                            sample_spatial_embedding_x=sample_spatial_embedding_x,
+                            sample_spatial_embedding_y=sample_spatial_embedding_y,
+                            sigma_x=sigma_x,
+                            sigma_y=sigma_y,
+                            color_sample=color_sample_dic,
+                            color_embedding=color_embedding_dic,
+                        )
                     else:
-                        visualizer.display(instances[0] > 0, key='center', title='Center', center_x=center_x,
-                                           center_y=center_y,
-                                           samples_x=samples_x, samples_y=samples_y,
-                                           sample_spatial_embedding_x=sample_spatial_embedding_x,
-                                           sample_spatial_embedding_y=sample_spatial_embedding_y,
-                                           sigma_x=sigma_x, sigma_y=sigma_y,
-                                           color_sample=color_sample_dic, color_embedding=color_embedding_dic)
-                visualizer.display(predictions.cpu(), key='prediction', title='Prediction')  # TODO
+                        visualizer.display(
+                            instances[0] > 0,
+                            key="center",
+                            title="Center",
+                            center_x=center_x,
+                            center_y=center_y,
+                            samples_x=samples_x,
+                            samples_y=samples_y,
+                            sample_spatial_embedding_x=sample_spatial_embedding_x,
+                            sample_spatial_embedding_y=sample_spatial_embedding_y,
+                            sigma_x=sigma_x,
+                            sigma_y=sigma_y,
+                            color_sample=color_sample_dic,
+                            color_embedding=color_embedding_dic,
+                        )
+                visualizer.display(
+                    predictions.cpu(), key="prediction", title="Prediction"
+                )  # TODO
 
     return loss_meter.avg
 
 
-def train_3d(virtual_batch_multiplier, one_hot, n_sigma, args):
+def train_3d(virtual_batch_multiplier, one_hot, n_sigma, args, device):
     """Trains 3D Model with virtual multiplier
 
-        Virtual batching code inspired from:
-        https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
-
-        Parameters
-        ----------
-        virtual_batch_multiplier : int
-            Set to 1 by default. Effective batch size is product of virtual_batch_multiplier and batch-size
-        one_hot : bool
-            In case the GT labels are available in one-hot fashion, this parameter is set equal to True
-            This parameter is not relevant for 3D data, and will be deprecated in a future code update
-        n_sigma: int
-            Should be equal to 3 for a 3D model
-        args: dictionary
-
-        Returns
-        -------
-        float
-            Average loss
-        """
+    Virtual batching code inspired from:
+    https://medium.com/huggingface/training-larger-batches-practical-tips
+    -on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
+
+    Parameters
+    ----------
+    virtual_batch_multiplier : int
+        Set to 1 by default. Effective batch size is product
+        of virtual_batch_multiplier and batch-size
+    one_hot : bool
+        In case the GT labels are available in one-hot fashion,
+        this parameter is set equal to True
+        This parameter is not relevant for 3D data, and will
+        be deprecated in a future code update
+    n_sigma: int
+        Should be equal to 3 for a 3D model
+    args: dictionary
+
+    device: torch.device
+
+    Returns
+    -------
+    float
+        Average loss
+    """
 
     # define meters
     loss_meter = AverageMeter()
     # put model into training mode
     model.train()
     for param_group in optimizer.param_groups:
-        print('learning rate: {}'.format(param_group['lr']))
+        print("learning rate: {}".format(param_group["lr"]))
 
     optimizer.zero_grad()  # Reset gradients tensors
     for i, sample in enumerate(tqdm(train_dataset_it)):
-        im = sample['image']
-        instances = sample['instance'].squeeze(1)
-        class_labels = sample['label'].squeeze(1)
-        center_images = sample['center_image'].squeeze(1)
+        im = sample["image"]
+        im = im.to(device)
+        instances = sample["instance"].squeeze(1).to(device)
+        class_labels = sample["label"].squeeze(1).to(device)
+        center_images = sample["center_image"].squeeze(1).to(device)
         output = model(im)  # Forward pass
 
         loss = criterion(output, instances, class_labels, center_images, **args)
         loss = loss / virtual_batch_multiplier  # Normalize our loss (if averaged)
         loss = loss.mean()
         loss.backward()  # Backward pass
-        if (i + 1) % virtual_batch_multiplier == 0:  # Wait for several backward steps
+        if (i + 1) % virtual_batch_multiplier == 0:
+            # Wait for several backward steps
             optimizer.step()  # Now we can do an optimizer step
             optimizer.zero_grad()  # Reset gradients tensors
         loss_meter.update(loss.item())
     return loss_meter.avg * virtual_batch_multiplier
 
 
-def train_vanilla_3d(display, display_embedding, display_it, one_hot, grid_x, grid_y, grid_z, pixel_x, pixel_y, pixel_z,
-                     n_sigma, zslice, args):
+def train_vanilla_3d(
+    display,
+    display_embedding,
+    display_it,
+    one_hot,
+    grid_x,
+    grid_y,
+    grid_z,
+    pixel_x,
+    pixel_y,
+    pixel_z,
+    n_sigma,
+    zslice,
+    args,
+    device,
+):
     """Trains 3D Model without virtual multiplier.
 
-            Parameters
-            ----------
-            display : bool
-                Displays input, GT, model predictions during training
-            display_embedding : bool
-                Displays embeddings for train (crop) images
-            display_it: int
-                Displays a new training image, the corresponding GT and model prediction every `display_it` crop images
-            one_hot: bool
-                In case the GT labels are available in one-hot fashion, this parameter is set equal to True
-                This parameter is not relevant for 3D data and will be deprecated in a future code update
-            grid_x: int
-                Number of pixels along x dimension which constitute a tile
-            grid_y: int
-                Number of pixels along y dimension which constitute a tile
-            grid_z: int
-                Number of pixels along z dimension which constitute a tile
-            pixel_x: float
-                The grid length along x is mapped to `pixel_x`.
-                For example, if grid_x = 1024 and pixel_x = 1.0, then each dX or the spacing between consecutive pixels
-                along the x dimension is set equal to pixel_x/grid_x = 1.0/1024
-            pixel_y: float
-                The grid length along y is mapped to `pixel_y`.
-                For example, if grid_y = 1024 and pixel_y = 1.0, then each dY or the spacing between consecutive pixels
-                along the y dimension is set equal to pixel_y/grid_y = 1.0/1024
-            pixel_z: float
-                The grid length along z is mapped to `pixel_z`.
-                For example, if grid_z = 1024 and pixel_z = 1.0, then each dY or the spacing between consecutive pixels
-                along the z dimension is set equal to pixel_z/grid_z = 1.0/1024
-            n_sigma: int
-                Should be set equal to 3 for a 3D model
-            zslice: int
-                If `display` = True, then the the raw image at z = z_slice is displayed during training
-            args: dictionary
-
-            Returns
-            -------
-            float
-                Average loss
-            """
+    Parameters
+    ----------
+    display : bool
+        Displays input, GT, model predictions during training
+    display_embedding : bool
+        Displays embeddings for train (crop) images
+    display_it: int
+        Displays a new training image, the corresponding GT
+        and model prediction every `display_it` crop images
+    one_hot: bool
+        In case the GT labels are available in one-hot fashion,
+        this parameter is set equal to True
+        This parameter is not relevant for 3D data
+        and will be deprecated in a future code update
+    grid_x: int
+        Number of pixels along x dimension which constitute a tile
+    grid_y: int
+        Number of pixels along y dimension which constitute a tile
+    grid_z: int
+        Number of pixels along z dimension which constitute a tile
+    pixel_x: float
+        The grid length along x is mapped to `pixel_x`.
+        For example, if grid_x = 1024 and pixel_x = 1.0,
+        then each dX or the spacing between consecutive pixels
+        along the x dimension is set equal to pixel_x/grid_x = 1.0/1024
+    pixel_y: float
+        The grid length along y is mapped to `pixel_y`.
+        For example, if grid_y = 1024 and pixel_y = 1.0,
+        then each dY or the spacing between consecutive pixels
+        along the y dimension is set equal to pixel_y/grid_y = 1.0/1024
+    pixel_z: float
+        The grid length along z is mapped to `pixel_z`.
+        For example, if grid_z = 1024 and pixel_z = 1.0,
+        then each dY or the spacing between consecutive pixels
+        along the z dimension is set equal to pixel_z/grid_z = 1.0/1024
+    n_sigma: int
+        Should be set equal to 3 for a 3D model
+    zslice: int
+        If `display` = True,
+        then the the raw image at z = z_slice is displayed during training
+    args: dictionary
+
+    device: torch.device
+
+    Returns
+    -------
+    float
+        Average loss
+    """
     # define meters
     loss_meter = AverageMeter()
     # put model into training mode
     model.train()
 
     for param_group in optimizer.param_groups:
-        print('learning rate: {}'.format(param_group['lr']))
+        print("learning rate: {}".format(param_group["lr"]))
 
     for i, sample in enumerate(tqdm(train_dataset_it)):
+        im = sample["image"].to(device)  # BCZYX
+        instances = sample["instance"].squeeze(1).to(device)  # BZYX
+        class_labels = sample["label"].squeeze(1).to(device)  # BZYX
+        center_images = sample["center_image"].squeeze(1).to(device)  # BZYX
 
-        im = sample['image']  # BCZYX
-        instances = sample['instance'].squeeze(1)  # BZYX
-        class_labels = sample['label'].squeeze(1)  # BZYX
-        center_images = sample['center_image'].squeeze(1)  # BZYX
         output = model(im)  # B 7 Z Y X
         loss = criterion(output, instances, class_labels, center_images, **args)
         loss = loss.mean()
@@ -276,43 +375,56 @@ def train_vanilla_3d(display, display_embedding, display_it, one_hot, grid_x, gr
         loss_meter.update(loss.item())
         if display and i % display_it == 0:
             with torch.no_grad():
-                visualizer.display(im[0, 0, zslice], key='image', title='Image')
-                predictions = cluster.cluster_with_gt(output[0], instances[0], n_sigma=n_sigma)
+                visualizer.display(im[0, 0, zslice], key="image", title="Image")
+                predictions = cluster.cluster_with_gt(
+                    output[0], instances[0], n_sigma=n_sigma
+                )
                 if one_hot:
                     instance = invert_one_hot(instances[0].cpu().detach().numpy())
-                    visualizer.display(instance, key='groundtruth', title='Ground Truth')  # TODO
+                    visualizer.display(
+                        instance, key="groundtruth", title="Ground Truth"
+                    )  # TODO
                     instance_ids = np.arange(instances.size(1))  # instances[0] --> DYX
                 else:
-                    visualizer.display(instances[0, zslice].cpu(), key='groundtruth', title='Ground Truth')  # TODO
+                    visualizer.display(
+                        instances[0, zslice].cpu(),
+                        key="groundtruth",
+                        title="Ground Truth",
+                    )  # TODO
                     instance_ids = instances[0].unique()
                     instance_ids = instance_ids[instance_ids != 0]
 
-                visualizer.display(predictions.cpu()[zslice, ...], key='prediction', title='Prediction')  # TODO
+                visualizer.display(
+                    predictions.cpu()[zslice, ...], key="prediction", title="Prediction"
+                )  # TODO
 
     return loss_meter.avg
 
 
-def val(virtual_batch_multiplier, one_hot, n_sigma, args):
+def val(virtual_batch_multiplier, one_hot, n_sigma, args, device):
     """Validates a 2D Model with virtual multiplier
 
-        Virtual batching code inspired from:
-        https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
-
-        Parameters
-        ----------
-        virtual_batch_multiplier : int
-            Set to 1 by default. Effective batch size is product of virtual_batch_multiplier and batch_size
-        one_hot : bool
-            In case the GT labels are available in one-hot fashion, this parameter is set equal to True
-        n_sigma: int
-            Should be equal to 2 for a 2D model
-        args: dictionary
-
-        Returns
-        -------
-        tuple: (float, float)
-            Average loss, Average IoU
-        """
+    Virtual batching code inspired from:
+    https://medium.com/huggingface/training-larger-batches-practical-
+    tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
+
+    Parameters
+    ----------
+    virtual_batch_multiplier : int
+        Set to 1 by default. Effective batch size is
+        product of virtual_batch_multiplier and batch_size
+    one_hot : bool
+        In case the GT labels are available in one-hot fashion,
+        this parameter is set equal to True
+    n_sigma: int
+        Should be equal to 2 for a 2D model
+    args: dictionary
+
+    Returns
+    -------
+    tuple: (float, float)
+        Average loss, Average IoU
+    """
 
     # define meters
     loss_meter, iou_meter = AverageMeter(), AverageMeter()
@@ -320,12 +432,20 @@ def val(virtual_batch_multiplier, one_hot, n_sigma, args):
     model.eval()
     with torch.no_grad():
         for i, sample in enumerate(tqdm(val_dataset_it)):
-            im = sample['image']
-            instances = sample['instance'].squeeze(1)
-            class_labels = sample['label'].squeeze(1)
-            center_images = sample['center_image'].squeeze(1)
+            im = sample["image"].to(device)
+            instances = sample["instance"].squeeze(1).to(device)
+            class_labels = sample["label"].squeeze(1).to(device)
+            center_images = sample["center_image"].squeeze(1).to(device)
             output = model(im)
-            loss = criterion(output, instances, class_labels, center_images, **args, iou=True, iou_meter=iou_meter)
+            loss = criterion(
+                output,
+                instances,
+                class_labels,
+                center_images,
+                **args,
+                iou=True,
+                iou_meter=iou_meter
+            )
             loss = loss.mean()
             loss = loss / virtual_batch_multiplier
             loss_meter.update(loss.item())
@@ -333,40 +453,56 @@ def val(virtual_batch_multiplier, one_hot, n_sigma, args):
     return loss_meter.avg * virtual_batch_multiplier, iou_meter.avg
 
 
-def val_vanilla(display, display_embedding, display_it, one_hot, grid_x, grid_y, pixel_x, pixel_y, n_sigma, args):
+def val_vanilla(
+    display,
+    display_embedding,
+    display_it,
+    one_hot,
+    grid_x,
+    grid_y,
+    pixel_x,
+    pixel_y,
+    n_sigma,
+    args,
+    device,
+):
     """Validates a 2D Model without virtual multiplier.
 
-            Parameters
-            ----------
-            display : bool
-                Displays input, GT, model predictions during training
-            display_embedding : bool
-                Displays embeddings for train (crop) images
-            display_it: int
-                Displays a new validation image, the corresponding GT and model prediction every `display_it` crop images
-            one_hot: bool
-                In case the GT labels are available in one-hot fashion, this parameter is set equal to True
-            grid_x: int
-                Number of pixels along x dimension which constitute a tile
-            grid_y: int
-                Number of pixels along y dimension which constitute a tile
-            pixel_x: float
-                The grid length along x is mapped to `pixel_x`.
-                For example, if grid_x = 1024 and pixel_x = 1.0, then each dX or the spacing between consecutive pixels
-                along the x dimension is set equal to pixel_x/grid_x = 1.0/1024
-            pixel_y: float
-                The grid length along y is mapped to `pixel_y`.
-                For example, if grid_y = 1024 and pixel_y = 1.0, then each dY or the spacing between consecutive pixels
-                along the y dimension is set equal to pixel_y/grid_y = 1.0/1024
-            n_sigma: int
-                Should be set equal to 2 for a 2D model
-            args: dictionary
-
-            Returns
-            -------
-            tuple: (float, float)
-                Average loss, Average IoU
-            """
+    Parameters
+    ----------
+    display : bool
+        Displays input, GT, model predictions during training
+    display_embedding : bool
+        Displays embeddings for train (crop) images
+    display_it: int
+        Displays a new validation image, the corresponding
+        GT and model prediction every `display_it` crop images
+    one_hot: bool
+        In case the GT labels are available in one-hot fashion,
+        this parameter is set equal to True
+    grid_x: int
+        Number of pixels along x dimension which constitute a tile
+    grid_y: int
+        Number of pixels along y dimension which constitute a tile
+    pixel_x: float
+        The grid length along x is mapped to `pixel_x`.
+        For example, if grid_x = 1024 and pixel_x = 1.0,
+        then each dX or the spacing between consecutive pixels
+        along the x dimension is set equal to pixel_x/grid_x = 1.0/1024
+    pixel_y: float
+        The grid length along y is mapped to `pixel_y`.
+        For example, if grid_y = 1024 and pixel_y = 1.0,
+        then each dY or the spacing between consecutive pixels
+        along the y dimension is set equal to pixel_y/grid_y = 1.0/1024
+    n_sigma: int
+        Should be set equal to 2 for a 2D model
+    args: dictionary
+
+    Returns
+    -------
+    tuple: (float, float)
+        Average loss, Average IoU
+    """
 
     # define meters
     loss_meter, iou_meter = AverageMeter(), AverageMeter()
@@ -374,93 +510,155 @@ def val_vanilla(display, display_embedding, display_it, one_hot, grid_x, grid_y,
     model.eval()
     with torch.no_grad():
         for i, sample in enumerate(tqdm(val_dataset_it)):
-            im = sample['image']
-            instances = sample['instance'].squeeze(1)
-            class_labels = sample['label'].squeeze(1)
-            center_images = sample['center_image'].squeeze(1)
+            im = sample["image"].to(device)
+            instances = sample["instance"].squeeze(1).to(device)
+            class_labels = sample["label"].squeeze(1).to(device)
+            center_images = sample["center_image"].squeeze(1).to(device)
             output = model(im)
-            loss = criterion(output, instances, class_labels, center_images, **args, iou=True, iou_meter=iou_meter)
+            loss = criterion(
+                output,
+                instances,
+                class_labels,
+                center_images,
+                **args,
+                iou=True,
+                iou_meter=iou_meter
+            )
             loss = loss.mean()
             if display and i % display_it == 0:
                 with torch.no_grad():
-                    visualizer.display(im[0], key='image', title='Image')
-                    predictions = cluster.cluster_with_gt(output[0], instances[0], n_sigma=n_sigma)
+                    visualizer.display(im[0], key="image", title="Image")
+                    predictions = cluster.cluster_with_gt(
+                        output[0], instances[0], n_sigma=n_sigma
+                    )
                     if one_hot:
                         instance = invert_one_hot(instances[0].cpu().detach().numpy())
-                        visualizer.display(instance, key='groundtruth', title='Ground Truth')  # TODO
+                        visualizer.display(
+                            instance, key="groundtruth", title="Ground Truth"
+                        )  # TODO
                         instance_ids = np.arange(instances.size(1))
                     else:
-                        visualizer.display(instances[0].cpu(), key='groundtruth', title='Ground Truth')  # TODO
+                        visualizer.display(
+                            instances[0].cpu(), key="groundtruth", title="Ground Truth"
+                        )  # TODO
                         instance_ids = instances[0].unique()
                         instance_ids = instance_ids[instance_ids != 0]
-                    if (display_embedding):
-                        center_x, center_y, samples_x, samples_y, sample_spatial_embedding_x, \
-                        sample_spatial_embedding_y, sigma_x, sigma_y, color_sample_dic, color_embedding_dic = \
-                            prepare_embedding_for_train_image(one_hot=one_hot, grid_x=grid_x, grid_y=grid_y,
-                                                              pixel_x=pixel_x, pixel_y=pixel_y,
-                                                              predictions=predictions, instance_ids=instance_ids,
-                                                              center_images=center_images,
-                                                              output=output, instances=instances, n_sigma=n_sigma)
+                    if display_embedding:
+                        (
+                            center_x,
+                            center_y,
+                            samples_x,
+                            samples_y,
+                            sample_spatial_embedding_x,
+                            sample_spatial_embedding_y,
+                            sigma_x,
+                            sigma_y,
+                            color_sample_dic,
+                            color_embedding_dic,
+                        ) = prepare_embedding_for_train_image(
+                            one_hot=one_hot,
+                            grid_x=grid_x,
+                            grid_y=grid_y,
+                            pixel_x=pixel_x,
+                            pixel_y=pixel_y,
+                            predictions=predictions,
+                            instance_ids=instance_ids,
+                            center_images=center_images,
+                            output=output,
+                            instances=instances,
+                            n_sigma=n_sigma,
+                        )
                         if one_hot:
-                            visualizer.display(torch.max(instances[0], dim=0)[0].cpu(), key='center', title='Center',
-                                               # torch.max returns a tuple
-                                               center_x=center_x,
-                                               center_y=center_y,
-                                               samples_x=samples_x, samples_y=samples_y,
-                                               sample_spatial_embedding_x=sample_spatial_embedding_x,
-                                               sample_spatial_embedding_y=sample_spatial_embedding_y,
-                                               sigma_x=sigma_x, sigma_y=sigma_y,
-                                               color_sample=color_sample_dic, color_embedding=color_embedding_dic)
+                            visualizer.display(
+                                torch.max(instances[0], dim=0)[0].cpu(),
+                                key="center",
+                                title="Center",
+                                # torch.max returns a tuple
+                                center_x=center_x,
+                                center_y=center_y,
+                                samples_x=samples_x,
+                                samples_y=samples_y,
+                                sample_spatial_embedding_x=sample_spatial_embedding_x,
+                                sample_spatial_embedding_y=sample_spatial_embedding_y,
+                                sigma_x=sigma_x,
+                                sigma_y=sigma_y,
+                                color_sample=color_sample_dic,
+                                color_embedding=color_embedding_dic,
+                            )
                         else:
-                            visualizer.display(instances[0] > 0, key='center', title='Center', center_x=center_x,
-                                               center_y=center_y,
-                                               samples_x=samples_x, samples_y=samples_y,
-                                               sample_spatial_embedding_x=sample_spatial_embedding_x,
-                                               sample_spatial_embedding_y=sample_spatial_embedding_y,
-                                               sigma_x=sigma_x, sigma_y=sigma_y,
-                                               color_sample=color_sample_dic, color_embedding=color_embedding_dic)
-
-                    visualizer.display(predictions.cpu(), key='prediction', title='Prediction')  # TODO
+                            visualizer.display(
+                                instances[0] > 0,
+                                key="center",
+                                title="Center",
+                                center_x=center_x,
+                                center_y=center_y,
+                                samples_x=samples_x,
+                                samples_y=samples_y,
+                                sample_spatial_embedding_x=sample_spatial_embedding_x,
+                                sample_spatial_embedding_y=sample_spatial_embedding_y,
+                                sigma_x=sigma_x,
+                                sigma_y=sigma_y,
+                                color_sample=color_sample_dic,
+                                color_embedding=color_embedding_dic,
+                            )
+
+                    visualizer.display(
+                        predictions.cpu(), key="prediction", title="Prediction"
+                    )  # TODO
 
             loss_meter.update(loss.item())
 
     return loss_meter.avg, iou_meter.avg
 
 
-def val_3d(virtual_batch_multiplier, one_hot, n_sigma, args):
+def val_3d(virtual_batch_multiplier, one_hot, n_sigma, args, device):
     """Validates a 3D Model with virtual multiplier
 
-            Virtual batching code inspired from:
-            https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
-
-            Parameters
-            ----------
-            virtual_batch_multiplier : int
-                Set to 1 by default. Effective batch size is product of virtual_batch_multiplier and batch-size
-            one_hot : bool
-                In case the GT labels are available in one-hot fashion, this parameter is set equal to True
-                This parameter is not relevant for 3D data, and will be deprecated in a future code update
-            n_sigma: int
-                Should be equal to 3 for a 3D model
-            args: dictionary
-
-            Returns
-            -------
-            tuple: (float, float)
-                Average loss, Average IoU
-            """
+    Virtual batching code inspired from:
+    https://medium.com/huggingface/training-larger-batches-practical-
+    tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
+
+    Parameters
+    ----------
+    virtual_batch_multiplier : int
+        Set to 1 by default. Effective batch size is
+        product of virtual_batch_multiplier and batch-size
+    one_hot : bool
+        In case the GT labels are available in one-hot fashion,
+        this parameter is set equal to True
+        This parameter is not relevant for 3D data,
+        and will be deprecated in a future code update
+    n_sigma: int
+        Should be equal to 3 for a 3D model
+    args: dictionary
+
+    device: torch.device
+
+    Returns
+    -------
+    tuple: (float, float)
+        Average loss, Average IoU
+    """
     # define meters
     loss_meter, iou_meter = AverageMeter(), AverageMeter()
     # put model into eval mode
     model.eval()
     with torch.no_grad():
         for i, sample in enumerate(tqdm(val_dataset_it)):
-            im = sample['image']
-            instances = sample['instance'].squeeze(1)
-            class_labels = sample['label'].squeeze(1)
-            center_images = sample['center_image'].squeeze(1)
+            im = sample["image"].to(device)
+            instances = sample["instance"].squeeze(1).to(device)
+            class_labels = sample["label"].squeeze(1).to(device)
+            center_images = sample["center_image"].squeeze(1).to(device)
             output = model(im)
-            loss = criterion(output, instances, class_labels, center_images, **args, iou=True, iou_meter=iou_meter)
+            loss = criterion(
+                output,
+                instances,
+                class_labels,
+                center_images,
+                **args,
+                iou=True,
+                iou_meter=iou_meter
+            )
             loss = loss.mean()
             loss = loss / virtual_batch_multiplier
             loss_meter.update(loss.item())
@@ -468,77 +666,120 @@ def val_3d(virtual_batch_multiplier, one_hot, n_sigma, args):
     return loss_meter.avg * virtual_batch_multiplier, iou_meter.avg
 
 
-def val_vanilla_3d(display, display_embedding, display_it, one_hot, grid_x, grid_y, grid_z, pixel_x, pixel_y, pixel_z,
-                   n_sigma, zslice, args):
+def val_vanilla_3d(
+    display,
+    display_embedding,
+    display_it,
+    one_hot,
+    grid_x,
+    grid_y,
+    grid_z,
+    pixel_x,
+    pixel_y,
+    pixel_z,
+    n_sigma,
+    zslice,
+    args,
+    device,
+):
     """Validates a 3D Model without virtual multiplier.
 
-                Parameters
-                ----------
-                display : bool
-                    Displays input, GT, model predictions during training
-                display_embedding : bool
-                    Displays embeddings for train (crop) images
-                display_it: int
-                    Displays a new training image, the corresponding GT and model prediction every `display_it` crop images
-                one_hot: bool
-                    In case the GT labels are available in one-hot fashion, this parameter is set equal to True
-                    This parameter is not relevant for 3D data and will be deprecated in a future code update
-                grid_x: int
-                    Number of pixels along x dimension which constitute a tile
-                grid_y: int
-                    Number of pixels along y dimension which constitute a tile
-                grid_z: int
-                    Number of pixels along z dimension which constitute a tile
-                pixel_x: float
-                    The grid length along x is mapped to `pixel_x`.
-                    For example, if grid_x = 1024 and pixel_x = 1.0, then each dX or the spacing between consecutive pixels
-                    along the x dimension is set equal to pixel_x/grid_x = 1.0/1024
-                pixel_y: float
-                    The grid length along y is mapped to `pixel_y`.
-                    For example, if grid_y = 1024 and pixel_y = 1.0, then each dY or the spacing between consecutive pixels
-                    along the y dimension is set equal to pixel_y/grid_y = 1.0/1024
-                pixel_z: float
-                    The grid length along z is mapped to `pixel_z`.
-                    For example, if grid_z = 1024 and pixel_z = 1.0, then each dY or the spacing between consecutive pixels
-                    along the z dimension is set equal to pixel_z/grid_z = 1.0/1024
-                n_sigma: int
-                    Should be set equal to 3 for a 3D model
-                zslice: int
-                    If `display` = True, then the the raw image at z = z_slice is displayed during training
-                args: dictionary
-
-                Returns
-                -------
-                tuple: (float, float)
-                    Average loss, Average IoU
-                """
+    Parameters
+    ----------
+    display : bool
+        Displays input, GT, model predictions during training
+    display_embedding : bool
+        Displays embeddings for train (crop) images
+    display_it: int
+        Displays a new training image, the corresponding GT
+        and model prediction every `display_it` crop images
+    one_hot: bool
+        In case the GT labels are available in one-hot fashion,
+        this parameter is set equal to True
+        This parameter is not relevant for 3D data
+        and will be deprecated in a future code update
+    grid_x: int
+        Number of pixels along x dimension which constitute a tile
+    grid_y: int
+        Number of pixels along y dimension which constitute a tile
+    grid_z: int
+        Number of pixels along z dimension which constitute a tile
+    pixel_x: float
+        The grid length along x is mapped to `pixel_x`.
+        For example, if grid_x = 1024 and pixel_x = 1.0, then each dX
+        or the spacing between consecutive pixels
+        along the x dimension is set equal to pixel_x/grid_x = 1.0/1024
+    pixel_y: float
+        The grid length along y is mapped to `pixel_y`.
+        For example, if grid_y = 1024 and pixel_y = 1.0, then each dY
+        or the spacing between consecutive pixels
+        along the y dimension is set equal to pixel_y/grid_y = 1.0/1024
+    pixel_z: float
+        The grid length along z is mapped to `pixel_z`.
+        For example, if grid_z = 1024 and pixel_z = 1.0, then each dY
+        or the spacing between consecutive pixels
+        along the z dimension is set equal to pixel_z/grid_z = 1.0/1024
+    n_sigma: int
+        Should be set equal to 3 for a 3D model
+    zslice: int
+        If `display` = True, then the the raw image at z = z_slice
+        is displayed during training
+    args: dictionary
+
+    device: torch.device
+
+    Returns
+    -------
+    tuple: (float, float)
+        Average loss, Average IoU
+    """
     # define meters
     loss_meter, iou_meter = AverageMeter(), AverageMeter()
     # put model into eval mode
     model.eval()
     with torch.no_grad():
         for i, sample in enumerate(tqdm(val_dataset_it)):
-            im = sample['image']  # BCZYX
-            instances = sample['instance'].squeeze(1)  # BZYX
-            class_labels = sample['label'].squeeze(1)  # BZYX
-            center_images = sample['center_image'].squeeze(1)  # BZYX
+            im = sample["image"].to(device)  # BCZYX
+            instances = sample["instance"].squeeze(1).to(device)  # BZYX
+            class_labels = sample["label"].squeeze(1).to(device)  # BZYX
+            center_images = sample["center_image"].squeeze(1).to(device)  # BZYX
             output = model(im)
-            loss = criterion(output, instances, class_labels, center_images, **args, iou=True, iou_meter=iou_meter)
+            loss = criterion(
+                output,
+                instances,
+                class_labels,
+                center_images,
+                **args,
+                iou=True,
+                iou_meter=iou_meter
+            )
             loss = loss.mean()
             if display and i % display_it == 0:
                 with torch.no_grad():
-                    visualizer.display(im[0, 0, zslice], key='image', title='Image')
-                    predictions = cluster.cluster_with_gt(output[0], instances[0], n_sigma=n_sigma)
+                    visualizer.display(im[0, 0, zslice], key="image", title="Image")
+                    predictions = cluster.cluster_with_gt(
+                        output[0], instances[0], n_sigma=n_sigma
+                    )
                     if one_hot:
                         instance = invert_one_hot(instances[0].cpu().detach().numpy())
-                        visualizer.display(instance, key='groundtruth', title='Ground Truth')  # TODO
+                        visualizer.display(
+                            instance, key="groundtruth", title="Ground Truth"
+                        )  # TODO
                         instance_ids = np.arange(instances.size(1))
                     else:
-                        visualizer.display(instances[0, zslice].cpu(), key='groundtruth', title='Ground Truth')  # TODO
+                        visualizer.display(
+                            instances[0, zslice].cpu(),
+                            key="groundtruth",
+                            title="Ground Truth",
+                        )  # TODO
                         instance_ids = instances[0].unique()
                         instance_ids = instance_ids[instance_ids != 0]
 
-                    visualizer.display(predictions.cpu()[zslice, ...], key='prediction', title='Prediction')  # TODO
+                    visualizer.display(
+                        predictions.cpu()[zslice, ...],
+                        key="prediction",
+                        title="Prediction",
+                    )  # TODO
 
             loss_meter.update(loss.item())
 
@@ -548,17 +789,18 @@ def val_vanilla_3d(display, display_embedding, display_it, one_hot, grid_x, grid
 def invert_one_hot(image):
     """Inverts a one-hot label mask.
 
-                Parameters
-                ----------
-                image : numpy array (I x H x W)
-                    Label mask present in one-hot fashion (i.e. with 0s and 1s and multiple z slices)
-                    here `I` is the number of GT or predicted objects
-
-                Returns
-                -------
-                numpy array (H x W)
-                    A flattened label mask with objects labelled from 1 ... I
-                """
+    Parameters
+    ----------
+    image : numpy array (I x H x W)
+        Label mask present in one-hot fashion
+        (i.e. with 0s and 1s and multiple z slices)
+        here `I` is the number of GT or predicted objects
+
+    Returns
+    -------
+    numpy array (H x W)
+        A flattened label mask with objects labelled from 1 ... I
+    """
     instance = np.zeros((image.shape[1], image.shape[2]), dtype="uint16")
     for z in range(image.shape[0]):
         instance = np.where(image[z] > 0, instance + z + 1, instance)
@@ -566,220 +808,336 @@ def invert_one_hot(image):
     return instance
 
 
-def save_checkpoint(state, is_best, epoch, save_dir, save_checkpoint_frequency, name='checkpoint.pth'):
+def save_checkpoint(
+    state, is_best, epoch, save_dir, save_checkpoint_frequency, name="checkpoint.pth"
+):
     """Trains 3D Model without virtual multiplier.
 
-                Parameters
-                ----------
-                state : dictionary
-                    The state of the model weights
-                is_best : bool
-                    In case the validation IoU is higher at the end of a certain epoch than previously recorded, `is_best` is set equal to True
-                epoch: int
-                    The current epoch
-                save_checkpoint_frequency: int
-                    The model weights are saved every `save_checkpoint_frequency` epochs
-                name: str, optional
-                    The model weights are saved under the name `name`
-
-                Returns
-                -------
-
-                """
-    print('=> saving checkpoint')
+    Parameters
+    ----------
+    state : dictionary
+        The state of the model weights
+    is_best : bool
+        In case the validation IoU is higher at the end of a certain epoch
+        than previously recorded, `is_best` is set equal to True
+    epoch: int
+        The current epoch
+    save_checkpoint_frequency: int
+        The model weights are saved every `save_checkpoint_frequency` epochs
+    name: str, optional
+        The model weights are saved under the name `name`
+
+    Returns
+    -------
+
+    """
+    print("=> saving checkpoint")
     file_name = os.path.join(save_dir, name)
     torch.save(state, file_name)
-    if (save_checkpoint_frequency is not None):
-        if (epoch % int(save_checkpoint_frequency) == 0):
+    if save_checkpoint_frequency is not None:
+        if epoch % int(save_checkpoint_frequency) == 0:
             file_name2 = os.path.join(save_dir, str(epoch) + "_" + name)
             torch.save(state, file_name2)
     if is_best:
-        shutil.copyfile(file_name, os.path.join(
-            save_dir, 'best_iou_model.pth'))
+        shutil.copyfile(file_name, os.path.join(save_dir, "best_iou_model.pth"))
 
 
-def begin_training(train_dataset_dict, val_dataset_dict, model_dict, loss_dict, configs, color_map='magma'):
+def begin_training(
+    train_dataset_dict,
+    val_dataset_dict,
+    model_dict,
+    loss_dict,
+    configs,
+    color_map="magma",
+):
     """Entry function for beginning the model training procedure.
 
-                Parameters
-                ----------
-                train_dataset_dict : dictionary
-                    Dictionary containing training data loader-specific parameters (for e.g. train_batch_size etc)
-                val_dataset_dict : dictionary
-                    Dictionary containing validation data loader-specific parameters (for e.g. val_batch_size etc)
-                model_dict: dictionary
-                    Dictionary containing model specific parameters (for e.g. number of outputs)
-                loss_dict: dictionary
-                    Dictionary containing loss specific parameters (for e.g. convex weights of different loss terms - w_iou, w_var etc)
-                configs: dictionary
-                    Dictionary containing general training parameters (for e.g. num_epochs, learning_rate etc)
-                color_map: str, optional
-                   Name of color map. Used in case configs['display'] is set equal to True
-
-                Returns
-                -------
-                """
-
-    if configs['save']:
-        if not os.path.exists(configs['save_dir']):
-            os.makedirs(configs['save_dir'])
-
-    if configs['display']:
+    Parameters
+    ----------
+    train_dataset_dict : dictionary
+        Dictionary containing training data loader-specific parameters
+        (for e.g. train_batch_size etc)
+    val_dataset_dict : dictionary
+        Dictionary containing validation data loader-specific parameters
+        (for e.g. val_batch_size etc)
+    model_dict: dictionary
+        Dictionary containing model specific parameters (for e.g. number of outputs)
+    loss_dict: dictionary
+        Dictionary containing loss specific parameters
+        (for e.g. convex weights of different loss terms - w_iou, w_var etc)
+    configs: dictionary
+        Dictionary containing general training parameters
+        (for e.g. num_epochs, learning_rate etc)
+    color_map: str, optional
+       Name of color map. Used in case configs['display'] is set equal to True
+
+    Returns
+    -------
+    """
+
+    if configs["save"]:
+        if not os.path.exists(configs["save_dir"]):
+            os.makedirs(configs["save_dir"])
+
+    if configs["display"]:
         plt.ion()
     else:
         plt.ioff()
         plt.switch_backend("agg")
 
     # set device
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = torch.device(configs["device"])
 
     # define global variables
     global train_dataset_it, val_dataset_it, model, criterion, optimizer, visualizer, cluster
 
     # train dataloader
 
-    train_dataset = get_dataset(train_dataset_dict['name'], train_dataset_dict['kwargs'])
-    train_dataset_it = torch.utils.data.DataLoader(train_dataset, batch_size=train_dataset_dict['batch_size'],
-                                                   shuffle=True, drop_last=True,
-                                                   num_workers=train_dataset_dict['workers'],
-                                                   pin_memory=True if configs['cuda'] else False)
+    train_dataset = get_dataset(
+        train_dataset_dict["name"], train_dataset_dict["kwargs"]
+    )
+    train_dataset_it = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_size=train_dataset_dict["batch_size"],
+        shuffle=True,
+        drop_last=True,
+        num_workers=train_dataset_dict["workers"],
+        pin_memory=True if configs["device"][:4] == "cuda" else False,
+    )
 
     # val dataloader
-    val_dataset = get_dataset(val_dataset_dict['name'], val_dataset_dict['kwargs'])
-    val_dataset_it = torch.utils.data.DataLoader(val_dataset, batch_size=val_dataset_dict['batch_size'], shuffle=True,
-                                                 drop_last=False, num_workers=val_dataset_dict['workers'],
-                                                 pin_memory=True if configs['cuda'] else False)
+    val_dataset = get_dataset(val_dataset_dict["name"], val_dataset_dict["kwargs"])
+    val_dataset_it = torch.utils.data.DataLoader(
+        val_dataset,
+        batch_size=val_dataset_dict["batch_size"],
+        shuffle=True,
+        drop_last=False,
+        num_workers=val_dataset_dict["workers"],
+        pin_memory=True if configs["device"][:4] == "cuda" else False,
+    )
 
     # set model
-    model = get_model(model_dict['name'], model_dict['kwargs'])
-    model.init_output(loss_dict['lossOpts']['n_sigma'])
-    model = torch.nn.DataParallel(model).to(device)
-
-    if (configs['grid_z'] is None):
-        criterion = get_loss(grid_z=None, grid_y=configs['grid_y'], grid_x=configs['grid_x'], pixel_z=None,
-                             pixel_y=configs['pixel_y'], pixel_x=configs['pixel_x'],
-                             one_hot=configs['one_hot'], loss_opts=loss_dict['lossOpts'])
+    model = get_model(model_dict["name"], model_dict["kwargs"])
+    model.init_output(loss_dict["lossOpts"]["n_sigma"])
+    model = model.to(device)
+    # model = torch.nn.DataParallel(model).to(device)
+
+    if configs["grid_z"] is None:
+        criterion = get_loss(
+            grid_z=None,
+            grid_y=configs["grid_y"],
+            grid_x=configs["grid_x"],
+            pixel_z=None,
+            pixel_y=configs["pixel_y"],
+            pixel_x=configs["pixel_x"],
+            one_hot=configs["one_hot"],
+            loss_opts=loss_dict["lossOpts"],
+        )
     else:
-        criterion = get_loss(configs['grid_z'], configs['grid_y'], configs['grid_x'],
-                             configs['pixel_z'], configs['pixel_y'], configs['pixel_x'],
-                             configs['one_hot'], loss_dict['lossOpts'])
-    criterion = torch.nn.DataParallel(criterion).to(device)
+        criterion = get_loss(
+            configs["grid_z"],
+            configs["grid_y"],
+            configs["grid_x"],
+            configs["pixel_z"],
+            configs["pixel_y"],
+            configs["pixel_x"],
+            configs["one_hot"],
+            loss_dict["lossOpts"],
+        )
+    criterion = criterion.to(device)
+    # criterion = torch.nn.DataParallel(criterion).to(device)
 
     # set optimizer
-    optimizer = torch.optim.Adam(model.parameters(), lr=configs['train_lr'], weight_decay=1e-4)
+    optimizer = torch.optim.Adam(
+        model.parameters(), lr=configs["train_lr"], weight_decay=1e-4
+    )
 
     def lambda_(epoch):
         return pow((1 - ((epoch) / 200)), 0.9)
 
-    #scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_, )
+    # scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_, )
 
-    if (configs['grid_z'] is None):
+    if configs["grid_z"] is None:
         # clustering
-        cluster = Cluster(configs['grid_y'], configs['grid_x'], configs['pixel_y'], configs['pixel_x'],
-                          configs['one_hot'])
+        cluster = Cluster(
+            configs["grid_y"],
+            configs["grid_x"],
+            configs["pixel_y"],
+            configs["pixel_x"],
+            device,
+            configs["one_hot"],
+        )
     else:
         # clustering
-        cluster = Cluster_3d(configs['grid_z'], configs['grid_y'], configs['grid_x'], configs['pixel_z'],
-                             configs['pixel_y'],
-                             configs['pixel_x'], configs['one_hot'])
+        cluster = Cluster_3d(
+            configs["grid_z"],
+            configs["grid_y"],
+            configs["grid_x"],
+            configs["pixel_z"],
+            configs["pixel_y"],
+            configs["pixel_x"],
+            device,
+            configs["one_hot"],
+        )
 
     # Visualizer
-    visualizer = Visualizer(('image', 'groundtruth', 'prediction', 'center'), color_map)  # 5 keys
+    visualizer = Visualizer(
+        ("image", "groundtruth", "prediction", "center"), color_map
+    )  # 5 keys
 
     # Logger
-    logger = Logger(('train', 'val', 'iou'), 'loss')
+    logger = Logger(("train", "val", "iou"), "loss")
 
     # resume
     start_epoch = 0
     best_iou = 0
-    if configs['resume_path'] is not None and os.path.exists(configs['resume_path']):
-        print('Resuming model from {}'.format(configs['resume_path']))
-        state = torch.load(configs['resume_path'])
-        start_epoch = state['epoch'] + 1
-        best_iou = state['best_iou']
-        model.load_state_dict(state['model_state_dict'], strict=True)
-        optimizer.load_state_dict(state['optim_state_dict'])
-        logger.data = state['logger_data']
-
-    for epoch in range(start_epoch, configs['n_epochs']):
-        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_, last_epoch=epoch - 1)
-        print('Starting epoch {}'.format(epoch))
+    if configs["resume_path"] is not None and os.path.exists(configs["resume_path"]):
+        print("Resuming model from {}".format(configs["resume_path"]))
+        state = torch.load(configs["resume_path"])
+        start_epoch = state["epoch"] + 1
+        best_iou = state["best_iou"]
+        model.load_state_dict(state["model_state_dict"], strict=True)
+        optimizer.load_state_dict(state["optim_state_dict"])
+        logger.data = state["logger_data"]
+
+    for epoch in range(start_epoch, configs["n_epochs"]):
+        scheduler = torch.optim.lr_scheduler.LambdaLR(
+            optimizer, lr_lambda=lambda_, last_epoch=epoch - 1
+        )
+        print("Starting epoch {}".format(epoch))
         # scheduler.step(epoch)
 
-        if (configs['grid_z'] is None):
-            if (train_dataset_dict['virtual_batch_multiplier'] > 1):
-                train_loss = train(virtual_batch_multiplier=train_dataset_dict['virtual_batch_multiplier'],
-                                   one_hot=configs['one_hot'],
-                                   n_sigma=loss_dict['lossOpts']['n_sigma'], args=loss_dict['lossW'])
-            elif (train_dataset_dict['virtual_batch_multiplier'] == 1):
-                train_loss = train_vanilla(display=configs['display'], display_embedding=configs['display_embedding'],
-                                           display_it=configs['display_it'], one_hot=configs['one_hot'],
-                                           n_sigma=loss_dict['lossOpts']['n_sigma'], grid_x=configs['grid_x'],
-                                           grid_y=configs['grid_y'],
-                                           pixel_x=configs['pixel_x'], pixel_y=configs['pixel_y'],
-                                           args=loss_dict['lossW'])
-
-            if (val_dataset_dict['virtual_batch_multiplier'] > 1):
-                val_loss, val_iou = val(virtual_batch_multiplier=val_dataset_dict['virtual_batch_multiplier'],
-                                        one_hot=configs['one_hot'],
-                                        n_sigma=loss_dict['lossOpts']['n_sigma'], args=loss_dict['lossW'])
-            elif (val_dataset_dict['virtual_batch_multiplier'] == 1):
-                val_loss, val_iou = val_vanilla(display=configs['display'],
-                                                display_embedding=configs['display_embedding'],
-                                                display_it=configs['display_it'], one_hot=configs['one_hot'],
-                                                n_sigma=loss_dict['lossOpts']['n_sigma'], grid_x=configs['grid_x'],
-                                                grid_y=configs['grid_y'], pixel_x=configs['pixel_x'],
-                                                pixel_y=configs['pixel_y'],
-                                                args=loss_dict['lossW'])
+        if configs["grid_z"] is None:
+            if train_dataset_dict["virtual_batch_multiplier"] > 1:
+                train_loss = train(
+                    virtual_batch_multiplier=train_dataset_dict[
+                        "virtual_batch_multiplier"
+                    ],
+                    one_hot=configs["one_hot"],
+                    n_sigma=loss_dict["lossOpts"]["n_sigma"],
+                    args=loss_dict["lossW"],
+                    device=device,
+                )
+            elif train_dataset_dict["virtual_batch_multiplier"] == 1:
+                train_loss = train_vanilla(
+                    display=configs["display"],
+                    display_embedding=configs["display_embedding"],
+                    display_it=configs["display_it"],
+                    one_hot=configs["one_hot"],
+                    n_sigma=loss_dict["lossOpts"]["n_sigma"],
+                    grid_x=configs["grid_x"],
+                    grid_y=configs["grid_y"],
+                    pixel_x=configs["pixel_x"],
+                    pixel_y=configs["pixel_y"],
+                    args=loss_dict["lossW"],
+                    device=device,
+                )
+
+            if val_dataset_dict["virtual_batch_multiplier"] > 1:
+                val_loss, val_iou = val(
+                    virtual_batch_multiplier=val_dataset_dict[
+                        "virtual_batch_multiplier"
+                    ],
+                    one_hot=configs["one_hot"],
+                    n_sigma=loss_dict["lossOpts"]["n_sigma"],
+                    args=loss_dict["lossW"],
+                    device=device,
+                )
+            elif val_dataset_dict["virtual_batch_multiplier"] == 1:
+                val_loss, val_iou = val_vanilla(
+                    display=configs["display"],
+                    display_embedding=configs["display_embedding"],
+                    display_it=configs["display_it"],
+                    one_hot=configs["one_hot"],
+                    n_sigma=loss_dict["lossOpts"]["n_sigma"],
+                    grid_x=configs["grid_x"],
+                    grid_y=configs["grid_y"],
+                    pixel_x=configs["pixel_x"],
+                    pixel_y=configs["pixel_y"],
+                    args=loss_dict["lossW"],
+                    device=device,
+                )
         else:
-            if (train_dataset_dict['virtual_batch_multiplier'] > 1):
-                train_loss = train_3d(virtual_batch_multiplier=train_dataset_dict['virtual_batch_multiplier'],
-                                      one_hot=configs['one_hot'],
-                                      n_sigma=loss_dict['lossOpts']['n_sigma'], args=loss_dict['lossW'])
-            elif (train_dataset_dict['virtual_batch_multiplier'] == 1):
-                train_loss = train_vanilla_3d(display=configs['display'],
-                                              display_embedding=configs['display_embedding'],
-                                              display_it=configs['display_it'], one_hot=configs['one_hot'],
-                                              n_sigma=loss_dict['lossOpts']['n_sigma'],
-                                              zslice=configs['display_zslice'], grid_x=configs['grid_x'],
-                                              grid_y=configs['grid_y'], grid_z=configs['grid_z'],
-                                              pixel_x=configs['pixel_x'], pixel_y=configs['pixel_y'],
-                                              pixel_z=configs['pixel_z'], args=loss_dict['lossW'], )
-
-            if (val_dataset_dict['virtual_batch_multiplier'] > 1):
-                val_loss, val_iou = val_3d(virtual_batch_multiplier=val_dataset_dict['virtual_batch_multiplier'],
-                                           one_hot=configs['one_hot'],
-                                           n_sigma=loss_dict['lossOpts']['n_sigma'], args=loss_dict['lossW'])
-            elif (val_dataset_dict['virtual_batch_multiplier'] == 1):
-                val_loss, val_iou = val_vanilla_3d(display=configs['display'],
-                                                   display_embedding=configs['display_embedding'],
-                                                   display_it=configs['display_it'], one_hot=configs['one_hot'],
-                                                   n_sigma=loss_dict['lossOpts']['n_sigma'],
-                                                   zslice=configs['display_zslice'], grid_x=configs['grid_x'],
-                                                   grid_y=configs['grid_y'], grid_z=configs['grid_z'],
-                                                   pixel_x=configs['pixel_x'], pixel_y=configs['pixel_y'],
-                                                   pixel_z=configs['pixel_z'],
-                                                   args=loss_dict['lossW'])
+            if train_dataset_dict["virtual_batch_multiplier"] > 1:
+                train_loss = train_3d(
+                    virtual_batch_multiplier=train_dataset_dict[
+                        "virtual_batch_multiplier"
+                    ],
+                    one_hot=configs["one_hot"],
+                    n_sigma=loss_dict["lossOpts"]["n_sigma"],
+                    args=loss_dict["lossW"],
+                    device=device,
+                )
+            elif train_dataset_dict["virtual_batch_multiplier"] == 1:
+                train_loss = train_vanilla_3d(
+                    display=configs["display"],
+                    display_embedding=configs["display_embedding"],
+                    display_it=configs["display_it"],
+                    one_hot=configs["one_hot"],
+                    n_sigma=loss_dict["lossOpts"]["n_sigma"],
+                    zslice=configs["display_zslice"],
+                    grid_x=configs["grid_x"],
+                    grid_y=configs["grid_y"],
+                    grid_z=configs["grid_z"],
+                    pixel_x=configs["pixel_x"],
+                    pixel_y=configs["pixel_y"],
+                    pixel_z=configs["pixel_z"],
+                    args=loss_dict["lossW"],
+                    device=device,
+                )
+
+            if val_dataset_dict["virtual_batch_multiplier"] > 1:
+                val_loss, val_iou = val_3d(
+                    virtual_batch_multiplier=val_dataset_dict[
+                        "virtual_batch_multiplier"
+                    ],
+                    one_hot=configs["one_hot"],
+                    n_sigma=loss_dict["lossOpts"]["n_sigma"],
+                    args=loss_dict["lossW"],
+                    device=device,
+                )
+            elif val_dataset_dict["virtual_batch_multiplier"] == 1:
+                val_loss, val_iou = val_vanilla_3d(
+                    display=configs["display"],
+                    display_embedding=configs["display_embedding"],
+                    display_it=configs["display_it"],
+                    one_hot=configs["one_hot"],
+                    n_sigma=loss_dict["lossOpts"]["n_sigma"],
+                    zslice=configs["display_zslice"],
+                    grid_x=configs["grid_x"],
+                    grid_y=configs["grid_y"],
+                    grid_z=configs["grid_z"],
+                    pixel_x=configs["pixel_x"],
+                    pixel_y=configs["pixel_y"],
+                    pixel_z=configs["pixel_z"],
+                    args=loss_dict["lossW"],
+                    device=device,
+                )
 
         scheduler.step()
-        print('===> train loss: {:.2f}'.format(train_loss))
-        print('===> val loss: {:.2f}, val iou: {:.2f}'.format(val_loss, val_iou))
+        print("===> train loss: {:.2f}".format(train_loss))
+        print("===> val loss: {:.2f}, val iou: {:.2f}".format(val_loss, val_iou))
 
-        logger.add('train', train_loss)
-        logger.add('val', val_loss)
-        logger.add('iou', val_iou)
-        logger.plot(save=configs['save'], save_dir=configs['save_dir'])  # TODO
+        logger.add("train", train_loss)
+        logger.add("val", val_loss)
+        logger.add("iou", val_iou)
+        logger.plot(save=configs["save"], save_dir=configs["save_dir"])  # TODO
 
         is_best = val_iou > best_iou
         best_iou = max(val_iou, best_iou)
 
-        if configs['save']:
+        if configs["save"]:
             state = {
-                'epoch': epoch,
-                'best_iou': best_iou,
-                'model_state_dict': model.state_dict(),
-                'optim_state_dict': optimizer.state_dict(),
-                'logger_data': logger.data,
+                "epoch": epoch,
+                "best_iou": best_iou,
+                "model_state_dict": model.state_dict(),
+                "optim_state_dict": optimizer.state_dict(),
+                "logger_data": logger.data,
             }
-        save_checkpoint(state, is_best, epoch, save_dir=configs['save_dir'],
-                        save_checkpoint_frequency=configs['save_checkpoint_frequency'])
+        save_checkpoint(
+            state,
+            is_best,
+            epoch,
+            save_dir=configs["save_dir"],
+            save_checkpoint_frequency=configs["save_checkpoint_frequency"],
+        )
diff --git a/EmbedSeg/utils/create_dicts.py b/EmbedSeg/utils/create_dicts.py
index 8329bbf..fbe27b6 100644
--- a/EmbedSeg/utils/create_dicts.py
+++ b/EmbedSeg/utils/create_dicts.py
@@ -4,245 +4,292 @@
 import EmbedSeg.utils.transforms as my_transforms
 
 
-def create_dataset_dict(data_dir,
-                        project_name,
-                        size,
-                        center,
-                        type,
-                        one_hot=False,
-                        name='2d',
-                        batch_size=16,
-                        virtual_batch_multiplier=1,
-                        workers=8,
-                        ):
+def create_dataset_dict(
+    data_dir,
+    project_name,
+    size,
+    center,
+    type,
+    one_hot=False,
+    name="2d",
+    batch_size=16,
+    virtual_batch_multiplier=1,
+    workers=8,
+):
     """
-        Creates `dataset_dict` dictionary from parameters.
-        Parameters
-        ----------
-        data_dir: string
-            Data is read from os.path.join(data_dir, project_name)
-        project_name: string
-            Data is read from os.path.join(data_dir, project_name)
-        size: int
-            Number of image-mask per epoch
-        center: string
-            One of 'centroid', 'approximate-medoid', 'medoid'
-        type: string
-            One of 'train', 'val'
-        one_hot: boolean
-            If 'True', instance images are perceived as DYX (here each object is encoded as 1 in its individual slice)
-            If 'False', instance image is perceived as YX and has the same dimensionality as raw image
-        name: string
-            One of '2d' or '3d'
-        batch_size: int
-            Effective Batch-size is the product of `batch_size` and `virtual_batch_multiplier`
-        virtual_batch_multiplier: int
-            Effective Batch-size is the product of `batch_size` and `virtual_batch_multiplier`
-        workers: int
-            Number of data-loader workers
+    Creates `dataset_dict` dictionary from parameters.
+    Parameters
+    ----------
+    data_dir: string
+        Data is read from os.path.join(data_dir, project_name)
+    project_name: string
+        Data is read from os.path.join(data_dir, project_name)
+    size: int
+        Number of image-mask per epoch
+    center: string
+        One of 'centroid', 'approximate-medoid', 'medoid'
+    type: string
+        One of 'train', 'val'
+    one_hot: boolean
+        If 'True', instance images are perceived as
+        DYX (here each object is encoded as 1 in its individual slice)
+        If 'False', instance image is perceived as
+        YX and has the same dimensionality as raw image
+    name: string
+        One of '2d' or '3d'
+    batch_size: int
+        Effective Batch-size is the product of `batch_size`
+        and `virtual_batch_multiplier`
+    virtual_batch_multiplier: int
+        Effective Batch-size is the product of `batch_size`
+        and `virtual_batch_multiplier`
+    workers: int
+        Number of data-loader workers
     """
-    if name == '2d':
-        set_transforms = my_transforms.get_transform([
-            {
-                'name': 'RandomRotationsAndFlips',
-                'opts': {
-                    'keys': ('image', 'instance', 'label', 'center_image'),
-                    'degrees': 90,
-                    'one_hot': one_hot,
-                }
-            },
-            {
-                'name': 'ToTensorFromNumpy',
-                'opts': {
-                    'keys': ('image', 'instance', 'label', 'center_image'),
-                    'type': (torch.FloatTensor, torch.ShortTensor, torch.ShortTensor, torch.BoolTensor),
-                }
-            },
-        ])
-    elif name == '3d':
-        set_transforms = my_transforms.get_transform([
-            {
-                'name': 'RandomRotationsAndFlips_3d',
-                'opts': {
-                    'keys': ('image', 'instance', 'label', 'center_image'),
-                    'degrees': 90,
-                    'one_hot': one_hot,
-                }
-            },
-            {
-                'name': 'ToTensorFromNumpy',
-                'opts': {
-                    'keys': ('image', 'instance', 'label', 'center_image'),
-                    'type': (torch.FloatTensor, torch.ShortTensor, torch.ShortTensor, torch.BoolTensor),
-                }
-            },
-        ])
+    if name == "2d":
+        set_transforms = my_transforms.get_transform(
+            [
+                {
+                    "name": "RandomRotationsAndFlips",
+                    "opts": {
+                        "keys": ("image", "instance", "label", "center_image"),
+                        "degrees": 90,
+                        "one_hot": one_hot,
+                    },
+                },
+                {
+                    "name": "ToTensorFromNumpy",
+                    "opts": {
+                        "keys": ("image", "instance", "label", "center_image"),
+                        "type": (
+                            torch.FloatTensor,
+                            torch.ShortTensor,
+                            torch.ShortTensor,
+                            torch.BoolTensor,
+                        ),
+                    },
+                },
+            ]
+        )
+    elif name == "3d":
+        set_transforms = my_transforms.get_transform(
+            [
+                {
+                    "name": "RandomRotationsAndFlips_3d",
+                    "opts": {
+                        "keys": ("image", "instance", "label", "center_image"),
+                        "degrees": 90,
+                        "one_hot": one_hot,
+                    },
+                },
+                {
+                    "name": "ToTensorFromNumpy",
+                    "opts": {
+                        "keys": ("image", "instance", "label", "center_image"),
+                        "type": (
+                            torch.FloatTensor,
+                            torch.ShortTensor,
+                            torch.ShortTensor,
+                            torch.BoolTensor,
+                        ),
+                    },
+                },
+            ]
+        )
     dataset_dict = {
-        'name': name,
-        'kwargs': {
-            'center': 'center-' + center,
-            'data_dir': os.path.join(data_dir, project_name),
-            'type': type,
-            'size': size,
-            'transform': set_transforms,
-            'one_hot': one_hot,
+        "name": name,
+        "kwargs": {
+            "center": "center-" + center,
+            "data_dir": os.path.join(data_dir, project_name),
+            "type": type,
+            "size": size,
+            "transform": set_transforms,
+            "one_hot": one_hot,
         },
-        'batch_size': batch_size,
-        'virtual_batch_multiplier': virtual_batch_multiplier,
-        'workers': workers,
-
+        "batch_size": batch_size,
+        "virtual_batch_multiplier": virtual_batch_multiplier,
+        "workers": workers,
     }
-    print("`{}_dataset_dict` dictionary successfully created with: \n -- {} images accessed from {}, "
-          "\n -- number of images per epoch equal to {}, "
-          "\n -- batch size set at {}, "
-          .format(type, type, os.path.join(data_dir, project_name, type, 'images'), size, batch_size,
-                  ))
+    print(
+        "`{}_dataset_dict` dictionary successfully created \
+                with: \n -- {} images accessed from {}, "
+        "\n -- number of images per epoch equal to {}, "
+        "\n -- batch size set at {}, ".format(
+            type,
+            type,
+            os.path.join(data_dir, project_name, type, "images"),
+            size,
+            batch_size,
+        )
+    )
     return dataset_dict
 
 
-def create_test_configs_dict(data_dir,
-                             checkpoint_path,
-                             norm,
-                             data_type,
-                             save_dir=None,
-                             normalization_factor=1.0,
-                             tta=True,
-                             one_hot=False,
-                             ap_val=0.5,
-                             seed_thresh=0.9,
-                             fg_thresh=0.5,
-                             min_object_size=10,
-                             mean_object_size=None,
-                             save_images=True,
-                             save_results=True,
-                             min_mask_sum=0,
-                             min_unclustered_sum=0,
-                             cuda=True,
-                             n_z=None,
-                             n_y=1024,
-                             n_x=1024,
-                             num_workers = 4,
-                             anisotropy_factor=None,
-                             l_y=1,
-                             l_x=1,
-                             name='2d',
-                             input_channels=1,
-                             type='test',
-                             normalization=True,
-                             cluster_fast = True,
-                             expand_grid=True,
-                             uniform_ds_factor =1,
-                             bg_id = 0
-                             ):
+def create_test_configs_dict(
+    data_dir,
+    checkpoint_path,
+    norm,
+    data_type,
+    save_dir=None,
+    normalization_factor=1.0,
+    tta=True,
+    one_hot=False,
+    ap_val=0.5,
+    seed_thresh=0.9,
+    fg_thresh=0.5,
+    min_object_size=10,
+    mean_object_size=None,
+    save_images=True,
+    save_results=True,
+    min_mask_sum=0,
+    min_unclustered_sum=0,
+    device="cuda:0",
+    n_z=None,
+    n_y=1024,
+    n_x=1024,
+    num_workers=4,
+    anisotropy_factor=None,
+    l_y=1,
+    l_x=1,
+    name="2d",
+    input_channels=1,
+    type="test",
+    normalization=True,
+    cluster_fast=True,
+    expand_grid=True,
+    uniform_ds_factor=1,
+    bg_id=0,
+):
     """
-        Creates `test_configs` dictionary from parameters.
-        Parameters
-        ----------
-        data_dir : str
-            Data is read from os.path.join(data_dir, 'test')
-        checkpoint_path: str
-            This indicates the path to the trained model
-        data_type: str
-            This reflects the data-type of the image and should be equal to one of '8-bit' or '16-bit'
-        save_dir: str
-            This indicates the directory where the results are saved
-        normalization_factor: float, default = 1.0
-            This, in general, does not need to be specified.
-            It is a legacy parameter, which will be deprecated in a later release.
-        tta: boolean
-            If True, then use test-time-augmentation
-            This parameter generally gives better results during inference, but increases the inference time.
-        one_hot: boolean
-            If True, then GT evaluation instance instance images are available in one-hot encoded style (DYX)
-            This parameter is not applicable for a 3D setting, since a pixel can be allocated to only one GT 3D object instance
-        ap_val: float, default = 0.5
-            While computing the AP_dsb, a true positive, false positive and false negative is determined
-            based on whether the IoU between the predicted and GT object >=ap_val
-        seed_thresh: float, default = 0.9
-            Only considers a subset of pixels for whom the seediness score is above the `seed_thresh`
-            These are then considered as potential object seeds
-        fg_thresh: float, default = 0.5
-            Only considers a subset of pixels for whom the seediness score is above the `fg_thresh`
-            These are then clustered into unique objects
-        min_object_size: int
-            Ignores objects having pixels less than min_object_size
-        mean_object_size: float
-            Average object size in terms of number of pixels
-            This is used in case running an Integer Linear Programming problem for combining objects
-        save_images: boolean
-            If True, then prediction images are saved
-        save_results: boolean
-            If True, then prediction results are saved in text file
-        min_mask_sum: int, default = 0
-            Only start creating instances, if there are at least `min_mask_sum` pixels in foreground!
-        min_unclustered_sum: int, default = 0
-            Stop when the number of seed candidates are less than `min_unclustered_sum`
-        cuda: boolean
-            True, indicates GPU usage
-        n_z: int
-            Size of the underlying grid along the z (or depth) dimension
-            Should be set to None, for a 2D setting
-        n_y: int
-            Size of the underlying grid along the y (or height) dimension
-        n_x: int
-            Size of the underlying grid along the x (or width) dimension
-        anisotropy_factor: float
-            Ratio of the sizes of the z and x pixel sizes
-            If the Raw Image is acquired in a down-sampled manner along the z dimension, then `anisotropy_factor` > 1
-        l_y: float
-            If l_y = 1.0 and n_y = 1024, then the pixel spacing along the y dimension is 1.0/1023
-        l_x: float
-            If l_x = 1.0 and n_x = 1024, then the pixel spacing along the x dimension is 1.0/1023
-        name: str
-            One of '2d', '3d', '3d_sliced'
-        input_channels: int
-            Number of channels in the Raw Image
-            For example, if Gray-scale, then 'input_channels' is equal to 1
-            If RGB, then 'input_channels' is equal to 3
-        type: str
-            One of 'train', 'val' or 'test'
-        normalization: bool
-            One of True or False
-            For test images, 'normalization' should be True
-            For other crops used during training, 'normalization' should be False since these image crops are already normalized
-        cluster_fast : bool, Default = True
-            If True, then both 'fg_thresh' and 'seed_thresh' are used
-            If False, then only 'fg_thresh' is used and seeds are identified as local maxima in the seediness map
-        expand_grid: bool, default = True
-            If True, then for a different sized test image, the grid and pixel spacing is automatically adjusted
-            If False, then prediction occurs on original grid or tile, and these predictions are then stitched
-        uniform_ds_factor: int, default = 1
-            If the original crops were generated by down-sampling in '01-data.ipynb', then the test images
-            should also be down-sampled before being input to the model and the predictions should then be up-sampled
-            to restore them back to the original size
-        bg_id : int, default = 0
-            Label of background in the ground truth test label masks
-            This parameter is only used while quantifying accuracy predicted label masks with ground truth label masks
-        num_workers: int, default = 4
+    Creates `test_configs` dictionary from parameters.
+    Parameters
+    ----------
+    data_dir : str
+        Data is read from os.path.join(data_dir, 'test')
+    checkpoint_path: str
+        This indicates the path to the trained model
+    data_type: str
+        This reflects the data-type of the image and should be equal
+        to one of '8-bit' or '16-bit'
+    save_dir: str
+        This indicates the directory where the results are saved
+    normalization_factor: float, default = 1.0
+        This, in general, does not need to be specified.
+        It is a legacy parameter, which will be deprecated in a later release.
+    tta: boolean
+        If True, then use test-time-augmentation
+        This parameter generally gives better results during inference,
+        but increases the inference time.
+    one_hot: boolean
+        If True, then GT evaluation instance instance images are available
+        in one-hot encoded style (DYX)
+        This parameter is not applicable for a 3D setting,
+        since a pixel can be allocated to only one GT 3D object instance
+    ap_val: float, default = 0.5
+        While computing the AP_dsb, a true positive, false positive and
+        false negative is determined
+        based on whether the IoU between the predicted and GT object >=ap_val
+    seed_thresh: float, default = 0.9
+        Only considers a subset of pixels for whom the seediness score
+        is above the `seed_thresh`
+        These are then considered as potential object seeds
+    fg_thresh: float, default = 0.5
+        Only considers a subset of pixels for whom
+        the seediness score is above the `fg_thresh`
+        These are then clustered into unique objects
+    min_object_size: int
+        Ignores objects having pixels less than min_object_size
+    mean_object_size: float
+        Average object size in terms of number of pixels
+        This is used in case running an Integer Linear Programming
+        problem for combining objects
+    save_images: boolean
+        If True, then prediction images are saved
+    save_results: boolean
+        If True, then prediction results are saved in text file
+    min_mask_sum: int, default = 0
+        Only start creating instances, if there are at least
+        `min_mask_sum` pixels in foreground!
+    min_unclustered_sum: int, default = 0
+        Stop when the number of seed candidates are less than `min_unclustered_sum`
+    device: string
+        Default: 'cuda:0'
+        Could set to 'mps' for acceleration with mac
+        Other options are 'cpu'
+    n_z: int
+        Size of the underlying grid along the z (or depth) dimension
+        Should be set to None, for a 2D setting
+    n_y: int
+        Size of the underlying grid along the y (or height) dimension
+    n_x: int
+        Size of the underlying grid along the x (or width) dimension
+    anisotropy_factor: float
+        Ratio of the sizes of the z and x pixel sizes
+        If the Raw Image is acquired in a down-sampled manner
+        along the z dimension, then `anisotropy_factor` > 1
+    l_y: float
+        If l_y = 1.0 and n_y = 1024, then the pixel spacing
+        along the y dimension is 1.0/1023
+    l_x: float
+        If l_x = 1.0 and n_x = 1024, then the pixel spacing
+        along the x dimension is 1.0/1023
+    name: str
+        One of '2d', '3d', '3d_sliced'
+    input_channels: int
+        Number of channels in the Raw Image
+        For example, if Gray-scale, then 'input_channels' is equal to 1
+        If RGB, then 'input_channels' is equal to 3
+    type: str
+        One of 'train', 'val' or 'test'
+    normalization: bool
+        One of True or False
+        For test images, 'normalization' should be True
+        For other crops used during training, 'normalization' should be False
+        since these image crops are already normalized
+    cluster_fast : bool, Default = True
+        If True, then both 'fg_thresh' and 'seed_thresh' are used
+        If False, then only 'fg_thresh' is used and seeds are identified as
+        local maxima in the seediness map
+    expand_grid: bool, default = True
+        If True, then for a different sized test image,
+        the grid and pixel spacing is automatically adjusted
+        If False, then prediction occurs on original grid or tile,
+        and these predictions are then stitched
+    uniform_ds_factor: int, default = 1
+        If the original crops were generated by down-sampling in '01-data.ipynb',
+        then the test images
+        should also be down-sampled before being input to the model
+        and the predictions should then be up-sampled
+        to restore them back to the original size
+    bg_id : int, default = 0
+        Label of background in the ground truth test label masks
+        This parameter is only used while quantifying accuracy predicted label
+        masks with ground truth label masks
+    num_workers: int, default = 4
     """
-    if (n_z is None):
+    if n_z is None:
         l_z = None
     else:
         l_z = (n_z - 1) / (n_x - 1) * anisotropy_factor
 
-    if name == '2d':
+    if name == "2d":
         n_sigma = 2
         num_classes = [4, 1]
-        model_name = 'branched_erfnet'
-    elif name == '3d':
+        model_name = "branched_erfnet"
+    elif name == "3d":
         n_sigma = 3
         num_classes = [6, 1]
-        model_name = 'branched_erfnet_3d'
-    elif name == '3d_sliced':
+        model_name = "branched_erfnet_3d"
+    elif name == "3d_sliced":
         n_sigma = 3
         num_classes = [4, 1]
-        model_name = 'branched_erfnet'
-    elif name =='3d_ilp':
+        model_name = "branched_erfnet"
+    elif name == "3d_ilp":
         n_sigma = 2
         num_classes = [4, 1]
-        model_name = 'branched_erfnet'
+        model_name = "branched_erfnet"
 
-    if name in ['3d_sliced']:
+    if name in ["3d_sliced"]:
         sliced_mode = True
     else:
         sliced_mode = False
@@ -257,7 +304,7 @@ def create_test_configs_dict(data_dir,
         tta=tta,
         seed_thresh=seed_thresh,
         fg_thresh=fg_thresh,
-        cuda=cuda,
+        device=device,
         save_results=save_results,
         save_images=save_images,
         save_dir=save_dir,
@@ -269,191 +316,217 @@ def create_test_configs_dict(data_dir,
         pixel_y=l_y,
         pixel_z=l_z,
         name=name,
-        num_workers = num_workers,
+        num_workers=num_workers,
         anisotropy_factor=anisotropy_factor,
-        cluster_fast = cluster_fast,
+        cluster_fast=cluster_fast,
         expand_grid=expand_grid,
-        uniform_ds_factor = uniform_ds_factor,
+        uniform_ds_factor=uniform_ds_factor,
         dataset={
-            'name': name,
-            'kwargs': {
-                'data_dir': data_dir,
-                'type': type,
-                'data_type': data_type,
-                'norm': norm,
-                'sliced_mode': sliced_mode,
-                'anisotropy_factor': anisotropy_factor,
-                'normalization': normalization,
-                'uniform_ds_factor': uniform_ds_factor,
-                'bg_id': bg_id,
-                'transform': my_transforms.get_transform([
-                    {
-                        'name': 'ToTensorFromNumpy',
-                        'opts': {
-                            'keys': ('image', 'instance', 'label'),
-                            'type': (torch.FloatTensor, torch.ShortTensor, torch.ShortTensor),
-                            'normalization_factor': normalization_factor
-                        }
-                    },
-                ]),
-                'one_hot': one_hot,
-            }
+            "name": name,
+            "kwargs": {
+                "data_dir": data_dir,
+                "type": type,
+                "data_type": data_type,
+                "norm": norm,
+                "sliced_mode": sliced_mode,
+                "anisotropy_factor": anisotropy_factor,
+                "normalization": normalization,
+                "uniform_ds_factor": uniform_ds_factor,
+                "bg_id": bg_id,
+                "transform": my_transforms.get_transform(
+                    [
+                        {
+                            "name": "ToTensorFromNumpy",
+                            "opts": {
+                                "keys": ("image", "instance", "label"),
+                                "type": (
+                                    torch.FloatTensor,
+                                    torch.ShortTensor,
+                                    torch.ShortTensor,
+                                ),
+                                "normalization_factor": normalization_factor,
+                            },
+                        },
+                    ]
+                ),
+                "one_hot": one_hot,
+            },
         },
-
         model={
-            'name': model_name,
-            'kwargs': {
-                'input_channels': input_channels,
-                'num_classes': num_classes,
-            }
-        }
+            "name": model_name,
+            "kwargs": {
+                "input_channels": input_channels,
+                "num_classes": num_classes,
+            },
+        },
     )
     print(
         "`test_configs` dictionary successfully created with: "
         "\n -- evaluation images accessed from {}, "
         "\n -- trained weights accessed from {}, "
         "\n -- output directory chosen as {}".format(
-            data_dir, checkpoint_path, save_dir))
+            data_dir, checkpoint_path, save_dir
+        )
+    )
     return test_configs
 
 
-def create_model_dict(input_channels, num_classes=[4, 1], name='2d'):
+def create_model_dict(input_channels, num_classes=[4, 1], name="2d"):
     """
-        Creates `model_dict` dictionary from parameters.
-        Parameters
-        ----------
-        input_channels: int
-            1 indicates gray-channle image, 3 indicates RGB image.
-        num_classes: list
-            [4, 1] -> 4 indicates offset in x, offset in y, margin in x, margin in y; 1 indicates seediness score
-        name: string
+    Creates `model_dict` dictionary from parameters.
+    Parameters
+    ----------
+    input_channels: int
+        1 indicates gray-channle image, 3 indicates RGB image.
+    num_classes: list
+        [4, 1] -> 4 indicates offset in x, offset in y,
+        margin in x, margin in y; 1 indicates seediness score
+    name: string
     """
     model_dict = {
-        'name': 'branched_erfnet' if name == '2d' else 'branched_erfnet_3d',
-        'kwargs': {
-            'num_classes': num_classes,
-            'input_channels': input_channels,
-        }
+        "name": "branched_erfnet" if name == "2d" else "branched_erfnet_3d",
+        "kwargs": {
+            "num_classes": num_classes,
+            "input_channels": input_channels,
+        },
     }
     print(
-        "`model_dict` dictionary successfully created with: \n -- num of classes equal to {}, \n -- input channels equal to {}, \n -- name equal to {}".format(
-            input_channels, num_classes, model_dict['name']))
+        "`model_dict` dictionary successfully created \
+                with: \n -- num of classes equal to {}, \n -- input channels \
+                equal to {}, \n -- name equal to {}".format(
+            input_channels, num_classes, model_dict["name"]
+        )
+    )
     return model_dict
 
 
 def create_loss_dict(foreground_weight=10, n_sigma=2, w_inst=1, w_var=10, w_seed=1):
     """
-        Creates `loss_dict` dictionary from parameters.
-        Parameters
-        ----------
-        foreground_weight: int
-        w_inst: int/float
-            weight on IOU loss
-        w_var: int/float
-            weight on variance loss
-        w_seed: int/float
-            weight on seediness loss
+    Creates `loss_dict` dictionary from parameters.
+    Parameters
+    ----------
+    foreground_weight: int
+    w_inst: int/float
+        weight on IOU loss
+    w_var: int/float
+        weight on variance loss
+    w_seed: int/float
+        weight on seediness loss
     """
-    loss_dict = {'lossOpts': {'n_sigma': n_sigma,
-                              'foreground_weight': foreground_weight,
-                              },
-                 'lossW': {'w_inst': w_inst,
-                           'w_var': w_var,
-                           'w_seed': w_seed,
-                           },
-                 }
+    loss_dict = {
+        "lossOpts": {
+            "n_sigma": n_sigma,
+            "foreground_weight": foreground_weight,
+        },
+        "lossW": {
+            "w_inst": w_inst,
+            "w_var": w_var,
+            "w_seed": w_seed,
+        },
+    }
     print(
-        "`loss_dict` dictionary successfully created with: \n -- foreground weight equal to {:.3f}, \n -- w_inst equal to {}, \n -- w_var equal to {}, \n -- w_seed equal to {}".format(
-            foreground_weight, w_inst, w_var, w_seed))
+        "`loss_dict` dictionary successfully created \
+                with: \n -- foreground weight equal to {:.3f}, \n -- w_inst \
+                equal to {}, \n -- w_var \
+                equal to {}, \n -- w_seed equal to {}".format(
+            foreground_weight, w_inst, w_var, w_seed
+        )
+    )
     return loss_dict
 
 
-def create_configs(save_dir,
-                   resume_path,
-                   one_hot=False,
-                   display=False,
-                   display_embedding=False,
-                   display_it=5,
-                   n_epochs=200,
-                   train_lr=5e-4,
-                   cuda=True,
-                   save=True,
-                   n_z=None,
-                   n_y=1024,
-                   n_x=1024,
-                   anisotropy_factor=None,
-                   l_y=1,
-                   l_x=1,
-                   save_checkpoint_frequency=None,
-                   display_zslice=None
-                   ):
+def create_configs(
+    save_dir,
+    resume_path,
+    one_hot=False,
+    display=False,
+    display_embedding=False,
+    display_it=5,
+    n_epochs=200,
+    train_lr=5e-4,
+    device="cuda:0",
+    save=True,
+    n_z=None,
+    n_y=1024,
+    n_x=1024,
+    anisotropy_factor=None,
+    l_y=1,
+    l_x=1,
+    save_checkpoint_frequency=None,
+    display_zslice=None,
+):
     """
-        Creates `configs` dictionary from parameters.
-        Parameters
-        ----------
-        save_dir: str
-            Path to where the experiment is saved
-        resume_path: str
-            Path to where the trained model (for example, checkpoint.pth) lives
-        one_hot: boolean
-            If 'True', instance images are perceived as DYX (here each object is encoded as 1 in its individual slice)
-            If 'False', instance image is perceived as YX and has the same dimensionality as raw image
-        display: boolean
-            If 'True', then realtime display of images, ground-truth, predictions are shown
-        display_embedding: boolean
-            If False, it suppresses embedding image
-        display_it: int
-            Shows display every n training/val steps (display_it = n)
-        n_epochs: int
-            Total number of epochs
-        train_lr: float
-            Starting learning rate
-        cuda: boolean
-            If True, use GPU
-        save: boolean
-            If True, then results are saved
-        grid_y: int
-            Size in y dimension of the largest evaluation image
-        grid_x: int
-            Size in x dimension of the largest evaluation image
-        pixel_y: int
-            Pixel size in y
-        pixel_x: int
-            Pixel size in x
-        save_checkpoint_frequency: int
-            Save model weights after 'n' epochs (in addition to last and best model weights)
-            Default is None
+    Creates `configs` dictionary from parameters.
+    Parameters
+    ----------
+    save_dir: str
+        Path to where the experiment is saved
+    resume_path: str
+        Path to where the trained model (for example, checkpoint.pth) lives
+    one_hot: boolean
+        If 'True', instance images are perceived as DYX
+        (here each object is encoded as 1 in its individual slice)
+        If 'False', instance image is perceived as YX and
+        has the same dimensionality as raw image
+    display: boolean
+        If 'True', then realtime display of images, ground-truth,
+        predictions are shown
+    display_embedding: boolean
+        If False, it suppresses embedding image
+    display_it: int
+        Shows display every n training/val steps (display_it = n)
+    n_epochs: int
+        Total number of epochs
+    train_lr: float
+        Starting learning rate
+    device: string
+        Default - 'cuda:0'
+        Set to 'mps' to train on Accellerated PyTorch training on Mac
+    save: boolean
+        If True, then results are saved
+    grid_y: int
+        Size in y dimension of the largest evaluation image
+    grid_x: int
+        Size in x dimension of the largest evaluation image
+    pixel_y: int
+        Pixel size in y
+    pixel_x: int
+        Pixel size in x
+    save_checkpoint_frequency: int
+        Save model weights after 'n' epochs (in addition to last and best model weights)
+        Default is None
     """
-    if (n_z is None):
+    if n_z is None:
         l_z = None
     else:
         l_z = (n_z - 1) / (n_x - 1) * anisotropy_factor
 
-    configs = dict(train_lr=train_lr,
-                   n_epochs=n_epochs,
-                   cuda=cuda,
-                   display=display,
-                   display_embedding=display_embedding,
-                   display_it=display_it,
-                   save=save,
-                   save_dir=save_dir,
-                   resume_path=resume_path,
-                   grid_z=n_z,
-                   grid_y=n_y,
-                   grid_x=n_x,
-                   pixel_z=l_z,
-                   pixel_y=l_y,
-                   pixel_x=l_x,
-                   one_hot=one_hot,
-                   save_checkpoint_frequency=save_checkpoint_frequency,
-                   display_zslice=display_zslice
-                   )
+    configs = dict(
+        train_lr=train_lr,
+        n_epochs=n_epochs,
+        device=device,
+        display=display,
+        display_embedding=display_embedding,
+        display_it=display_it,
+        save=save,
+        save_dir=save_dir,
+        resume_path=resume_path,
+        grid_z=n_z,
+        grid_y=n_y,
+        grid_x=n_x,
+        pixel_z=l_z,
+        pixel_y=l_y,
+        pixel_x=l_x,
+        one_hot=one_hot,
+        save_checkpoint_frequency=save_checkpoint_frequency,
+        display_zslice=display_zslice,
+    )
     print(
         "`configs` dictionary successfully created with: "
         "\n -- n_epochs equal to {}, "
         "\n -- save_dir equal to {}, "
         "\n -- n_z equal to {}, "
         "\n -- n_y equal to {}, "
-        "\n -- n_x equal to {}, "
-        .format(n_epochs, save_dir, n_z, n_y, n_x))
+        "\n -- n_x equal to {}, ".format(n_epochs, save_dir, n_z, n_y, n_x)
+    )
     return configs
diff --git a/EmbedSeg/utils/generate_crops.py b/EmbedSeg/utils/generate_crops.py
index fa6e7a8..80ca86e 100644
--- a/EmbedSeg/utils/generate_crops.py
+++ b/EmbedSeg/utils/generate_crops.py
@@ -13,24 +13,26 @@
 
 def _fill_label_holes(lbl_img, **kwargs):
     """
-        Function taken from Stardist repository: https://github.com/stardist/stardist
+    Function taken from Stardist repository: https://github.com/stardist/stardist
     """
     lbl_img_filled = np.zeros_like(lbl_img)
-    for l in (set(np.unique(lbl_img)) - set([0])):
-        mask = lbl_img == l
+    for lbl in set(np.unique(lbl_img)) - set([0]):
+        mask = lbl_img == lbl
         mask_filled = binary_fill_holes(mask, **kwargs)
-        lbl_img_filled[mask_filled] = l
+        lbl_img_filled[mask_filled] = lbl
     return lbl_img_filled
 
 
 def fill_label_holes(lbl_img, **kwargs):
     """
-        Fill small holes in label image.
-        Function taken from Stardist repository: https://github.com/stardist/stardist
+    Fill small holes in label image.
+    Function taken from Stardist repository: https://github.com/stardist/stardist
     """
 
     def grow(sl, interior):
-        return tuple(slice(s.start - int(w[0]), s.stop + int(w[1])) for s, w in zip(sl, interior))
+        return tuple(
+            slice(s.start - int(w[0]), s.stop + int(w[1])) for s, w in zip(sl, interior)
+        )
 
     def shrink(interior):
         return tuple(slice(int(w[0]), (-1 if w[1] else None)) for w in interior)
@@ -38,7 +40,8 @@ def shrink(interior):
     objects = find_objects(lbl_img)
     lbl_img_filled = np.zeros_like(lbl_img)
     for i, sl in enumerate(objects, 1):
-        if sl is None: continue
+        if sl is None:
+            continue
         interior = [(s.start > 0, s.stop < sz) for s, sz in zip(sl, lbl_img.shape)]
         shrink_slice = shrink(interior)
         grown_mask = lbl_img[grow(sl, interior)] == i
@@ -47,20 +50,22 @@ def shrink(interior):
     return lbl_img_filled
 
 
-def normalize_min_max_percentile(x, pmin=3, pmax=99.8, axis=None, clip=False, eps=1e-20, dtype=np.float32):
+def normalize_min_max_percentile(
+    x, pmin=3, pmax=99.8, axis=None, clip=False, eps=1e-20, dtype=float
+):
     """
-        Percentile-based image normalization.
-        Function taken from StarDist repository  https://github.com/stardist/stardist
+    Percentile-based image normalization.
+    Function taken from StarDist repository  https://github.com/stardist/stardist
     """
     mi = np.percentile(x, pmin, axis=axis, keepdims=True)
     ma = np.percentile(x, pmax, axis=axis, keepdims=True)
     return normalize_mi_ma(x, mi, ma, clip=clip, eps=eps, dtype=dtype)
 
 
-def normalize_mi_ma(x, mi, ma, clip=False, eps=1e-20, dtype=np.float32):
+def normalize_mi_ma(x, mi, ma, clip=False, eps=1e-20, dtype=float):
     """
-        Percentile-based image normalization.
-        Function taken from StarDist repository  https://github.com/stardist/stardist
+    Percentile-based image normalization.
+    Function taken from StarDist repository  https://github.com/stardist/stardist
     """
     if dtype is not None:
         x = x.astype(dtype, copy=False)
@@ -70,6 +75,7 @@ def normalize_mi_ma(x, mi, ma, clip=False, eps=1e-20, dtype=np.float32):
 
     try:
         import numexpr
+
         x = numexpr.evaluate("(x - mi) / ( ma - mi + eps )")
     except ImportError:
         x = (x - mi) / (ma - mi + eps)
@@ -89,13 +95,14 @@ def normalize_mean_std(x, axis=None):
         Input Image
     axis: int
         Set equal to 0, 1, 2 ... in case you normalize based on dimension
-        Can be useful if the image contains multiple channels and each channel should be independently normalized
+        Can be useful if the image contains multiple channels and each channel
+        should be independently normalized
     Returns
     -------
     numpy array: float
         Normalized image with mean equal to 0 and standard deviation equal to 1
     """
-    x = x.astype(np.float32)
+    x = x.astype(float)
     if axis is None:
         mean = np.mean(x)
         std = np.std(x)
@@ -109,20 +116,22 @@ def normalize_mean_std(x, axis=None):
 
 @jit(nopython=True)
 def pairwise_python(X):
-    """Helper function to compute pairwise Euclidean distance for a matrix of size M x N containing M,  N-dimensional row entries
+    """Helper function to compute pairwise Euclidean distance for a matrix
+    of size M x N containing M,  N-dimensional row entries
 
-        Parameters
-        ----------
-        x : numpy array
-            Input Image
-        Returns
-        -------
-        D : numpy array, M x M
-            D[i, j] corresponds to the Euclidean Distance between the i th and j th rows of X
-        """
+    Parameters
+    ----------
+    x : numpy array
+        Input Image
+    Returns
+    -------
+    D : numpy array, M x M
+        D[i, j] corresponds to the Euclidean Distance between
+        the i th and j th rows of X
+    """
     M = X.shape[0]
     N = X.shape[1]
-    D = np.empty((M, M), dtype=np.float32)
+    D = np.empty((M, M), dtype=float)
     for i in range(M):
         for j in range(M):
             d = 0.0
@@ -135,43 +144,46 @@ def pairwise_python(X):
 
 def generate_center_image(instance, center, ids, one_hot):
     """
-        Generates a `center_image` which is one (True) for all center locations and zero (False) otherwise.
-
-        Parameters
-        ----------
-        instance: numpy array
-            `instance` image containing unique `ids` for each object (YX)
-             or present in a one-hot encoded style where each object is one in it own slice and zero elsewhere.
-        center: string
-            One of 'centroid', 'approximate-medoid' or 'medoid'.
-        ids: list
-            Unique ids corresponding to the objects present in the instance image.
-        one_hot: boolean
-            True (in this case, `instance` has shape DYX) or False (in this case, `instance` has shape YX).
-
-        Returns
-        -------
-        numpy array: bool
-        Center image with center locations set to True
+    Generates a `center_image` which is one (True) for all center locations
+    and zero (False) otherwise.
+
+    Parameters
+    ----------
+    instance: numpy array
+        `instance` image containing unique `ids` for each object (YX)
+         or present in a one-hot encoded style where each object is one
+         in its own slice and zero elsewhere.
+    center: string
+        One of 'centroid', 'approximate-medoid' or 'medoid'.
+    ids: list
+        Unique ids corresponding to the objects present in the instance image.
+    one_hot: boolean
+        True (in this case, `instance` has shape DYX) or False
+        (in this case, `instance` has shape YX).
+
+    Returns
+    -------
+    numpy array: bool
+    Center image with center locations set to True
     """
 
-    if (not one_hot):
+    if not one_hot:
         center_image = np.zeros(instance.shape, dtype=bool)
     else:
         center_image = np.zeros((instance.shape[-2], instance.shape[-1]), dtype=bool)
     for j, id in enumerate(ids):
-        if (not one_hot):
+        if not one_hot:
             y, x = np.where(instance == id)
         else:
             y, x = np.where(instance[id] == 1)
         if len(y) != 0 and len(x) != 0:
-            if (center == 'centroid'):
+            if center == "centroid":
                 ym, xm = np.mean(y), np.mean(x)
-            elif (center == 'approximate-medoid'):
+            elif center == "approximate-medoid":
                 ym_temp, xm_temp = np.median(y), np.median(x)
                 imin = np.argmin((x - xm_temp) ** 2 + (y - ym_temp) ** 2)
                 ym, xm = y[imin], x[imin]
-            elif (center == 'medoid'):
+            elif center == "medoid":
                 dist_matrix = pairwise_python(np.vstack((x, y)).transpose())
                 imin = np.argmin(np.sum(dist_matrix, axis=0))
                 ym, xm = y[imin], x[imin]
@@ -179,66 +191,85 @@ def generate_center_image(instance, center, ids, one_hot):
     return center_image
 
 
-def generate_center_image_3d(instance, center, ids, one_hot, anisotropy_factor, speed_up):
+def generate_center_image_3d(
+    instance, center, ids, one_hot, anisotropy_factor, speed_up
+):
+    """
+    Generates a `center_image` for 3D image crops which is one (True)
+    for all center locations and zero (False) otherwise.
+
+    Parameters
+    ----------
+    instance: numpy array
+        `instance` image containing unique `ids` for each object (ZYX)
+    center: string
+        One of 'centroid', 'approximate-medoid' or 'medoid'.
+    ids: list
+        Unique ids corresponding to the objects present in the instance image.
+    one_hot: boolean
+        This parameter is not used in the 3D setting and will be deprecated.
+    speed_up: int
+        This computes the centers of crops faster by down-sampling crops
+        along x and y dimensions.
+
+    Returns
+    -------
+    numpy array: bool
+    Center image with center locations set to True
     """
-        Generates a `center_image` for 3D image crops which is one (True) for all center locations and zero (False) otherwise.
-
-        Parameters
-        ----------
-        instance: numpy array
-            `instance` image containing unique `ids` for each object (ZYX)
-        center: string
-            One of 'centroid', 'approximate-medoid' or 'medoid'.
-        ids: list
-            Unique ids corresponding to the objects present in the instance image.
-        one_hot: boolean
-            This parameter is not used in the 3D setting and will be deprecated.
-        speed_up: int
-            This computes the centers of crops faster by down-sampling crops along x and y dimensions.
-
-        Returns
-        -------
-        numpy array: bool
-        Center image with center locations set to True
-        """
 
     center_image = np.zeros(instance.shape, dtype=bool)
-    instance_downsampled = instance[:, ::int(speed_up), ::int(speed_up)]  # down sample in x and y
+    instance_downsampled = instance[
+        :, :: int(speed_up), :: int(speed_up)
+    ]  # down sample in x and y
     for j, id in enumerate(ids):
         z, y, x = np.where(instance_downsampled == id)
         if len(y) != 0 and len(x) != 0:
-            if (center == 'centroid'):
+            if center == "centroid":
                 zm, ym, xm = np.mean(z), np.mean(y), np.mean(x)
-            elif (center == 'approximate-medoid'):
+            elif center == "approximate-medoid":
                 zm_temp, ym_temp, xm_temp = np.median(z), np.median(y), np.median(x)
-                imin = np.argmin((x - xm_temp) ** 2 + (y - ym_temp) ** 2 + (anisotropy_factor * (z - zm_temp)) ** 2)
+                imin = np.argmin(
+                    (x - xm_temp) ** 2
+                    + (y - ym_temp) ** 2
+                    + (anisotropy_factor * (z - zm_temp)) ** 2
+                )
                 zm, ym, xm = z[imin], y[imin], x[imin]
-            elif (center == 'medoid'):
+            elif center == "medoid":
                 dist_matrix = pairwise_python(
-                    np.vstack((speed_up * x, speed_up * y, anisotropy_factor * z)).transpose())
+                    np.vstack(
+                        (speed_up * x, speed_up * y, anisotropy_factor * z)
+                    ).transpose()
+                )
                 imin = np.argmin(np.sum(dist_matrix, axis=0))
                 zm, ym, xm = z[imin], y[imin], x[imin]
-            center_image[int(np.round(zm)), int(np.round(speed_up * ym)), int(np.round(speed_up * xm))] = True
+            center_image[
+                int(np.round(zm)),
+                int(np.round(speed_up * ym)),
+                int(np.round(speed_up * xm)),
+            ] = True
     return center_image
 
 
 def sparsify(instance):
     """
-        Ensures that no slice in the `instance` image crop is empty (completely zero)
+    Ensures that no slice in the `instance` image crop is empty (completely zero)
 
-        Parameters
-        ----------
-        instance: numpy array,
-            `instance` image present in a one-hot encoded style where each object is one in it own slice and zero elsewhere.
+    Parameters
+    ----------
+    instance: numpy array,
+        `instance` image present in a one-hot encoded style where each object
+        is one in it own slice and zero elsewhere.
 
-        Returns
-        ----------
-        numpy array
-        Numpy array with number of slices which are fewer or equal to the number of slices in `instance`
+    Returns
+    ----------
+    numpy array
+    Numpy array with number of slices which are fewer or equal to
+    the number of slices in `instance`
     """
     instance_sparse = []
     for z in range(instance.shape[0]):
-        if (np.sum(instance[z, ...]) > 0):
+        if np.sum(instance[z, ...]) > 0:
             instance_sparse.append(instance[z, ...])
     return np.array(instance_sparse)
 
@@ -246,75 +277,91 @@ def sparsify(instance):
 def encode(filename, img, one_hot=False):
     """Convert label mask (img) in a run-length encoded fashion and save to a csv file
 
-        Parameters
-        ----------
-        filename : str
-            Name of csv file
-        img: numpy array
-            Label mask image containing instance segmentations
-        one_hot: boolean, optional
-            If the Label mask is present in a one-hot encoded fashion, set this equal to True
-        Returns
-        -------
-        """
+    Parameters
+    ----------
+    filename : str
+        Name of csv file
+    img: numpy array
+        Label mask image containing instance segmentations
+    one_hot: boolean, optional
+        If the Label mask is present in a one-hot encoded fashion,
+        set this equal to True
+    Returns
+    -------
+    """
     data = []
 
     if one_hot:
         ids = np.arange(img.shape[0])
         for id in ids:
             d = rletools.encode(np.asfortranarray(img[id]) == 1)
-            data.append([id, d['counts'], d['size']])
+            data.append([id, d["counts"], d["size"]])
     else:
         ids = np.unique(img)
         ids = ids[ids != 0]
         for id in ids:
             d = rletools.encode(np.asfortranarray(img) == id)
-            data.append([id, d['counts'], d['size']])
+            data.append([id, d["counts"], d["size"]])
     df = pd.DataFrame(data)
     df.to_csv(filename, index=False, header=None)
 
 
-def process(im, inst, crops_dir, data_subset, crop_size, center, norm='min-max-percentile', one_hot=False,
-            data_type='8-bit', normalization_factor=None,
-            rle_encode=False, fraction_max_ids=1.0, background_id=0):
+def process(
+    im,
+    inst,
+    crops_dir,
+    data_subset,
+    crop_size,
+    center,
+    norm="min-max-percentile",
+    one_hot=False,
+    data_type="8-bit",
+    normalization_factor=None,
+    rle_encode=False,
+    fraction_max_ids=1.0,
+    background_id=0,
+):
     """Entry function which generates 2D crops from 2D images (one_hot is False)
 
-            Parameters
-            ----------
-            im: numpy array
-                raw image
-            inst: numpy array
-                label mask
-            crops_dir: str
-                Path to where the crops are saved
-            data_subset: str
-                Set equal to one of `train` or `val`
-            crop_size: int
-                Height and Width of crop
-            center: str
-                Set equal to one of `centroid`, `medoid` or `approximate-medoid`
-            norm: str
-                Set equal to one of `min-max-percentile` or `absolute` or `mean-std`
-            one_hot: bool
-                If label masks are available in one-hot encoded fashion, set `one_hot` equal to True
-            data_type: str
-                Set equal to `8-bit` or `16-bit`
-            normalization_factor: int
-                In case, norm is set equal to `absolute`, then the image intensities are divided by the `normalization_factor`
-            rle_encode: bool
-                If set equal to True, the label masks are saved as csv files
-            fraction_max_ids: float, between 0 and 1.0
-                If set equal to a value less than 1.0, then only that fraction of ids are processed are used to make crops
-            background_id: int, optional
-                Id of the background in the label mask
-
-            Returns
-            -------
-            """
-
-    image_path = os.path.join(crops_dir, data_subset, 'images/')
-    instance_path = os.path.join(crops_dir, data_subset, 'masks/')
-    center_image_path = os.path.join(crops_dir, data_subset, 'center-' + center + '/')
+    Parameters
+    ----------
+    im: numpy array
+        raw image
+    inst: numpy array
+        label mask
+    crops_dir: str
+        Path to where the crops are saved
+    data_subset: str
+        Set equal to one of `train` or `val`
+    crop_size: int
+        Height and Width of crop
+    center: str
+        Set equal to one of `centroid`, `medoid` or `approximate-medoid`
+    norm: str
+        Set equal to one of `min-max-percentile` or `absolute` or `mean-std`
+    one_hot: bool
+        If label masks are available in one-hot
+        encoded fashion, set `one_hot` equal to True
+    data_type: str
+        Set equal to `8-bit` or `16-bit`
+    normalization_factor: int
+        In case, norm is set equal to `absolute`, then the image intensities
+        are divided by the `normalization_factor`
+    rle_encode: bool
+        If set equal to True, the label masks are saved as csv files
+    fraction_max_ids: float, between 0 and 1.0
+        If set equal to a value less than 1.0, then only that fraction of ids
+        are processed are used to make crops
+    background_id: int, optional
+        Id of the background in the label mask
+
+    Returns
+    -------
+    """
+
+    image_path = os.path.join(crops_dir, data_subset, "images/")
+    instance_path = os.path.join(crops_dir, data_subset, "masks/")
+    center_image_path = os.path.join(crops_dir, data_subset, "center-" + center + "/")
 
     if not os.path.exists(image_path):
         os.makedirs(os.path.dirname(image_path))
@@ -327,25 +374,25 @@ def process(im, inst, crops_dir, data_subset, crop_size, center, norm='min-max-p
         print("Created new directory : {}".format(center_image_path))
 
     instance = tifffile.imread(inst).astype(np.uint16)
-    image = tifffile.imread(im).astype(np.float32)
+    image = tifffile.imread(im).astype(float)
 
-    if (norm == 'min-max-percentile'):
+    if norm == "min-max-percentile":
         if image.ndim == 2:  # gray-scale
             image = normalize_min_max_percentile(image, 1, 99.8, axis=(0, 1))
         elif image.ndim == 3:  # multi-channel image (C, H, W)
             image = normalize_min_max_percentile(image, 1, 99.8, axis=(1, 2))
-    elif (norm == 'mean-std'):
+    elif norm == "mean-std":
         if image.ndim == 2:
             image = normalize_mean_std(image)  # axis == None
         elif image.ndim == 3:
             image = normalize_mean_std(image, axis=(1, 2))
-    elif (norm == 'absolute'):
-        if data_type == '8-bit':
+    elif norm == "absolute":
+        if data_type == "8-bit":
             if normalization_factor is None:
                 image /= 255
             else:
                 image /= normalization_factor
-        elif data_type == '16-bit':
+        elif data_type == "16-bit":
             if normalization_factor is None:
                 image /= 65535
             else:
@@ -374,90 +421,165 @@ def process(im, inst, crops_dir, data_subset, crop_size, center, norm='min-max-p
         jj = int(np.clip(ym - crop_size / 2, 0, h - crop_size))
         ii = int(np.clip(xm - crop_size / 2, 0, w - crop_size))
         if image.ndim == 2:
-            if (image[jj:jj + crop_size, ii:ii + crop_size].shape == (crop_size, crop_size)):
-                im_crop = image[jj:jj + crop_size, ii:ii + crop_size]
-                instance_crop = instance_np[jj:jj + crop_size, ii:ii + crop_size]
-                center_image_crop = generate_center_image(instance_crop, center, ids, one_hot)
-                tifffile.imsave(image_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j), im_crop)
+            if image[jj : jj + crop_size, ii : ii + crop_size].shape == (
+                crop_size,
+                crop_size,
+            ):
+                im_crop = image[jj : jj + crop_size, ii : ii + crop_size]
+                instance_crop = instance_np[jj : jj + crop_size, ii : ii + crop_size]
+                center_image_crop = generate_center_image(
+                    instance_crop, center, ids, one_hot
+                )
+                tifffile.imsave(
+                    image_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
+                    im_crop,
+                )
                 if rle_encode:
-                    encode(instance_path + os.path.basename(im)[:-4] + "_{:03d}.csv".format(j), instance_crop)
-                    encode(center_image_path + os.path.basename(im)[:-4] + "_{:03d}.csv".format(j),
-                           center_image_crop)
+                    encode(
+                        instance_path
+                        + os.path.basename(im)[:-4]
+                        + "_{:03d}.csv".format(j),
+                        instance_crop,
+                    )
+                    encode(
+                        center_image_path
+                        + os.path.basename(im)[:-4]
+                        + "_{:03d}.csv".format(j),
+                        center_image_crop,
+                    )
                 else:
-                    tifffile.imsave(instance_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
-                                    instance_crop)
-                    tifffile.imsave(center_image_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
-                                    center_image_crop)
+                    tifffile.imsave(
+                        instance_path
+                        + os.path.basename(im)[:-4]
+                        + "_{:03d}.tif".format(j),
+                        instance_crop,
+                    )
+                    tifffile.imsave(
+                        center_image_path
+                        + os.path.basename(im)[:-4]
+                        + "_{:03d}.tif".format(j),
+                        center_image_crop,
+                    )
         elif image.ndim == 3:
-            if (image[:, jj:jj + crop_size, ii:ii + crop_size].shape == (c, crop_size, crop_size)):
-                im_crop = image[:, jj:jj + crop_size, ii:ii + crop_size]
-                instance_crop = instance[jj:jj + crop_size, ii:ii + crop_size]
-                center_image_crop = generate_center_image(instance_crop, center, ids, one_hot)
-                tifffile.imsave(image_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j), im_crop)
+            if image[:, jj : jj + crop_size, ii : ii + crop_size].shape == (
+                c,
+                crop_size,
+                crop_size,
+            ):
+                im_crop = image[:, jj : jj + crop_size, ii : ii + crop_size]
+                instance_crop = instance[jj : jj + crop_size, ii : ii + crop_size]
+                center_image_crop = generate_center_image(
+                    instance_crop, center, ids, one_hot
+                )
+                tifffile.imsave(
+                    image_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
+                    im_crop,
+                )
                 if rle_encode:
-                    encode(instance_path + os.path.basename(im)[:-4] + "_{:03d}.csv".format(j), instance_crop)
-                    encode(center_image_path + os.path.basename(im)[:-4] + "_{:03d}.csv".format(j), center_image_crop)
+                    encode(
+                        instance_path
+                        + os.path.basename(im)[:-4]
+                        + "_{:03d}.csv".format(j),
+                        instance_crop,
+                    )
+                    encode(
+                        center_image_path
+                        + os.path.basename(im)[:-4]
+                        + "_{:03d}.csv".format(j),
+                        center_image_crop,
+                    )
                 else:
-                    tifffile.imsave(instance_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
-                                    instance_crop)
-                    tifffile.imsave(center_image_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
-                                    center_image_crop)
+                    tifffile.imsave(
+                        instance_path
+                        + os.path.basename(im)[:-4]
+                        + "_{:03d}.tif".format(j),
+                        instance_crop,
+                    )
+                    tifffile.imsave(
+                        center_image_path
+                        + os.path.basename(im)[:-4]
+                        + "_{:03d}.tif".format(j),
+                        center_image_crop,
+                    )
+
+
+def process_3d(
+    im,
+    inst,
+    crops_dir,
+    data_subset,
+    crop_size_x,
+    crop_size_y,
+    crop_size_z,
+    center,
+    norm="min-max-percentile",
+    one_hot=False,
+    anisotropy_factor=1.0,
+    speed_up=1.0,
+    data_type="8-bit",
+    normalization_factor=None,
+    rle_encode=False,
+    fraction_max_ids=1.0,
+    background_id=0,
+    uniform_ds_factor=1,
+):
+    """Entry function which generates 3D crops from 3D images
 
+    Parameters
+    ----------
+    im: numpy array
+        raw image
+    inst: numpy array
+        label mask
+    crops_dir: str
+        Path to where the crops are saved
+    data_subset: str
+        Set equal to one of `train` or `val`
+    crop_size_x: int
+        Width of crop
+    crop_size_y: int
+        Height of crop
+    crop_size_z: int
+        Depth of crop
+    center: str
+        Set equal to one of `centroid`, `medoid` or `approximate-medoid`
+    norm: str
+        Set equal to one of `min-max-percentile` or `absolute` or `mean-std`
+    one_hot: bool
+        If label masks are available in one-hot encoded fashion,
+        set `one_hot` equal to True
+        This parameter should be always set equal to False for 3D processing
+        and will be deprecated in a later release
+    anisotropy_factor: float
+        This parameter should be set equal to the ratio of the z pixel size
+        to the x or y pixel size
+        Here, we assume that the x or y pixel size is the same
+        (If the imaging is down-sampled in the z dimension,
+         `anisotropy_factor` is greater than 1.0)
+    data_type: str
+        Set equal to `8-bit` or `16-bit`
+    normalization_factor: int
+        In case, norm is set equal to `absolute`,
+        then the image intensities are divided by the `normalization_factor`
+    rle_encode: bool
+        If set equal to True, the label masks are saved as csv files
+    fraction_max_ids: float, between 0 and 1.0
+        If set equal to a value less than 1.0, then only that fraction of ids
+        are processed are used to make crops
+    background_id: int, optional
+        Id of the background in the label
+    uniform_ds_factor: int, optional
+        In case, the image and corresponding GT instance should be down-sampled
+        This serves the purpose of increasing the receptive field
+        without increasing the GPU memory requirement
 
-def process_3d(im, inst, crops_dir, data_subset, crop_size_x, crop_size_y, crop_size_z, center,
-               norm='min-max-percentile',
-               one_hot=False, anisotropy_factor=1.0, speed_up=1.0, data_type='8-bit', normalization_factor=None,
-               rle_encode=False, fraction_max_ids=1.0, background_id=0, uniform_ds_factor = 1):
-    """Entry function which generates 3D crops from 3D images
+    Returns
+    -------
+    """
 
-                Parameters
-                ----------
-                im: numpy array
-                    raw image
-                inst: numpy array
-                    label mask
-                crops_dir: str
-                    Path to where the crops are saved
-                data_subset: str
-                    Set equal to one of `train` or `val`
-                crop_size_x: int
-                    Width of crop
-                crop_size_y: int
-                    Height of crop
-                crop_size_z: int
-                    Depth of crop
-                center: str
-                    Set equal to one of `centroid`, `medoid` or `approximate-medoid`
-                norm: str
-                    Set equal to one of `min-max-percentile` or `absolute` or `mean-std`
-                one_hot: bool
-                    If label masks are available in one-hot encoded fashion, set `one_hot` equal to True
-                    This parameter should be always set equal to False for 3D processing and will be deprecated in a later release
-                anisotropy_factor: float
-                    This parameter should be set equal to the ratio of the z pixel size to the x or y pixel size
-                    Here, we assume that the x or y pixel size is the same
-                    (If the imaging is down-sampled in the z dimension, `anisotropy_factor` is greater than 1.0)
-                data_type: str
-                    Set equal to `8-bit` or `16-bit`
-                normalization_factor: int
-                    In case, norm is set equal to `absolute`, then the image intensities are divided by the `normalization_factor`
-                rle_encode: bool
-                    If set equal to True, the label masks are saved as csv files
-                fraction_max_ids: float, between 0 and 1.0
-                    If set equal to a value less than 1.0, then only that fraction of ids are processed are used to make crops
-                background_id: int, optional
-                    Id of the background in the label
-                uniform_ds_factor: int, optional
-                    In case, the image and corresponding GT instance should be down-sampled
-                    This serves the purpose of increasing the receptive field without increasing the GPU memory requirement
-
-                Returns
-                -------
-                """
-
-    image_path = os.path.join(crops_dir, data_subset, 'images/')
-    instance_path = os.path.join(crops_dir, data_subset, 'masks/')
-    center_image_path = os.path.join(crops_dir, data_subset, 'center-' + center + '/')
+    image_path = os.path.join(crops_dir, data_subset, "images/")
+    instance_path = os.path.join(crops_dir, data_subset, "masks/")
+    center_image_path = os.path.join(crops_dir, data_subset, "center-" + center + "/")
 
     if not os.path.exists(image_path):
         os.makedirs(os.path.dirname(image_path))
@@ -470,39 +592,45 @@ def process_3d(im, inst, crops_dir, data_subset, crop_size_x, crop_size_y, crop_
         print("Created new directory : {}".format(center_image_path))
 
     instance = tifffile.imread(inst).astype(np.uint16)
-    image = tifffile.imread(im).astype(np.float32)
-
-
+    image = tifffile.imread(im).astype(float)
 
-    if (norm == 'min-max-percentile'):
+    if norm == "min-max-percentile":
         image = normalize_min_max_percentile(image, 1, 99.8, axis=(0, 1, 2))
-    elif (norm == 'mean-std'):
+    elif norm == "mean-std":
         image = normalize_mean_std(image)
-    elif (norm == 'absolute'):
-        if data_type == '8-bit':
+    elif norm == "absolute":
+        if data_type == "8-bit":
             if normalization_factor is None:
                 image /= 255
             else:
                 image /= normalization_factor
-        elif data_type == '16-bit':
+        elif data_type == "16-bit":
             if normalization_factor is None:
                 image /= 65535
             else:
                 image /= normalization_factor
     instance = fill_label_holes(instance)
 
-    # sometimes it helps to downsample the image and instance, in order to increase the receptive field
-    instance_ds = instance[::uniform_ds_factor, ::uniform_ds_factor, ::uniform_ds_factor]
+    # sometimes it helps to downsample the image and instance,
+    # in order to increase the receptive field
+    instance_ds = instance[
+        ::uniform_ds_factor, ::uniform_ds_factor, ::uniform_ds_factor
+    ]
     image_ds = image[::uniform_ds_factor, ::uniform_ds_factor, ::uniform_ds_factor]
 
-    # but we would still ike the downsampled image to be of the same size as the original image, which we ensure through padding
+    # but we would still like the downsampled image to be of the same
+    # size as the original image, which we ensure through padding
     dz = image.shape[0] - image_ds.shape[0]
     dy = image.shape[1] - image_ds.shape[1]
     dx = image.shape[2] - image_ds.shape[2]
-    pad_width = ((dz//2, dz-dz//2),(dy//2, dy-dy//2), (dx//2, dx-dx//2))
+    pad_width = (
+        (dz // 2, dz - dz // 2),
+        (dy // 2, dy - dy // 2),
+        (dx // 2, dx - dx // 2),
+    )
 
-    image= np.pad(image_ds, pad_width)
-    instance = np.pad(instance_ds, pad_width, 'constant', constant_values=background_id)
+    image = np.pad(image_ds, pad_width)
+    instance = np.pad(instance_ds, pad_width, "constant", constant_values=background_id)
 
     d, h, w = image.shape
     instance_np = np.array(instance, copy=False)
@@ -520,61 +648,103 @@ def process_3d(im, inst, crops_dir, data_subset, crop_size_x, crop_size_y, crop_
         jj = int(np.clip(ym - crop_size_y / 2, 0, h - crop_size_y))
         ii = int(np.clip(xm - crop_size_x / 2, 0, w - crop_size_x))
 
-        if (image[kk:kk + crop_size_z, jj:jj + crop_size_y, ii:ii + crop_size_x].shape == (
-                crop_size_z, crop_size_y, crop_size_x)):
-            im_crop = image[kk:kk + crop_size_z, jj:jj + crop_size_y, ii:ii + crop_size_x]
-            instance_crop = instance_np[kk:kk + crop_size_z, jj:jj + crop_size_y, ii:ii + crop_size_x]
-            center_image_crop = generate_center_image_3d(instance_crop, center, ids, one_hot=one_hot,
-                                                         anisotropy_factor=anisotropy_factor, speed_up=speed_up)
-            tifffile.imsave(image_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j), im_crop)
+        if image[
+            kk : kk + crop_size_z, jj : jj + crop_size_y, ii : ii + crop_size_x
+        ].shape == (crop_size_z, crop_size_y, crop_size_x):
+            im_crop = image[
+                kk : kk + crop_size_z, jj : jj + crop_size_y, ii : ii + crop_size_x
+            ]
+            instance_crop = instance_np[
+                kk : kk + crop_size_z, jj : jj + crop_size_y, ii : ii + crop_size_x
+            ]
+            center_image_crop = generate_center_image_3d(
+                instance_crop,
+                center,
+                ids,
+                one_hot=one_hot,
+                anisotropy_factor=anisotropy_factor,
+                speed_up=speed_up,
+            )
+            tifffile.imsave(
+                image_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
+                im_crop,
+            )
             if rle_encode:
-                encode(instance_path + os.path.basename(im)[:-4] + "_{:03d}.csv".format(j),
-                       instance_crop.astype(np.uint16))
-                encode(center_image_path + os.path.basename(im)[:-4] + "_{:03d}.csv".format(j), center_image_crop)
+                encode(
+                    instance_path + os.path.basename(im)[:-4] + "_{:03d}.csv".format(j),
+                    instance_crop.astype(np.uint16),
+                )
+                encode(
+                    center_image_path
+                    + os.path.basename(im)[:-4]
+                    + "_{:03d}.csv".format(j),
+                    center_image_crop,
+                )
             else:
-                tifffile.imsave(instance_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
-                                instance_crop.astype(np.uint16))
-                tifffile.imsave(center_image_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
-                                center_image_crop)
-
-
-def process_one_hot(im, inst, crops_dir, data_subset, crop_size, center, one_hot=True, norm='min-max-percentile',
-                    data_type='8-bit', normalization_factor=None, fraction_max_ids=1.0, rle_encode=False):
+                tifffile.imsave(
+                    instance_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
+                    instance_crop.astype(np.uint16),
+                )
+                tifffile.imsave(
+                    center_image_path
+                    + os.path.basename(im)[:-4]
+                    + "_{:03d}.tif".format(j),
+                    center_image_crop,
+                )
+
+
+def process_one_hot(
+    im,
+    inst,
+    crops_dir,
+    data_subset,
+    crop_size,
+    center,
+    one_hot=True,
+    norm="min-max-percentile",
+    data_type="8-bit",
+    normalization_factor=None,
+    fraction_max_ids=1.0,
+    rle_encode=False,
+):
     """Entry function which generates 2D crops from 2D images (one_hot is True)
 
-                Parameters
-                ----------
-                im: numpy array
-                    raw image
-                inst: numpy array
-                    label mask
-                crops_dir: str
-                    Path to where the crops are saved
-                data_subset: str
-                    Set equal to one of `train` or `val`
-                crop_size: int
-                    Height and Width of crop
-                center: str
-                    Set equal to one of `centroid`, `medoid` or `approximate-medoid`
-                one_hot: bool
-                    If label masks are available in one-hot encoded fashion, set `one_hot` equal to True
-                norm: str
-                    Set equal to one of `min-max-percentile` or `absolute` or `mean-std`
-                data_type: str
-                    Set equal to `8-bit` or `16-bit`
-                normalization_factor: int
-                    In case, norm is set equal to `absolute`, then the image intensities are divided by the `normalization_factor`
-                fraction_max_ids: float, between 0 and 1.0
-                    If set equal to a value less than 1.0, then only that fraction of ids are processed are used to make crops
-                rle_encode: bool
-                    If set equal to True, the label masks are saved as csv files
-                Returns
-                -------
-                """
-
-    image_path = os.path.join(crops_dir, data_subset, 'images/')
-    instance_path = os.path.join(crops_dir, data_subset, 'masks/')
-    center_image_path = os.path.join(crops_dir, data_subset, 'center-' + center + '/')
+    Parameters
+    ----------
+    im: numpy array
+        raw image
+    inst: numpy array
+        label mask
+    crops_dir: str
+        Path to where the crops are saved
+    data_subset: str
+        Set equal to one of `train` or `val`
+    crop_size: int
+        Height and Width of crop
+    center: str
+        Set equal to one of `centroid`, `medoid` or `approximate-medoid`
+    one_hot: bool
+        If label masks are available in one-hot encoded fashion,
+        set `one_hot` equal to True
+    norm: str
+        Set equal to one of `min-max-percentile` or `absolute` or `mean-std`
+    data_type: str
+        Set equal to `8-bit` or `16-bit`
+    normalization_factor: int
+        In case, norm is set equal to `absolute`,
+        then the image intensities are divided by the `normalization_factor`
+    fraction_max_ids: float, between 0 and 1.0
+        If set equal to a value less than 1.0,
+        then only that fraction of ids are processed are used to make crops
+    rle_encode: bool
+        If set equal to True, the label masks are saved as csv files
+    Returns
+    -------
+    """
+
+    image_path = os.path.join(crops_dir, data_subset, "images/")
+    instance_path = os.path.join(crops_dir, data_subset, "masks/")
+    center_image_path = os.path.join(crops_dir, data_subset, "center-" + center + "/")
 
     try:
         os.makedirs(os.path.dirname(image_path))
@@ -584,19 +754,19 @@ def process_one_hot(im, inst, crops_dir, data_subset, crop_size, center, one_hot
         pass
 
     instance = tifffile.imread(inst).astype(np.uint16)
-    image = tifffile.imread(im).astype(np.float32)
+    image = tifffile.imread(im).astype(float)
 
-    if (norm == 'min-max-percentile'):
+    if norm == "min-max-percentile":
         image = normalize_min_max_percentile(image, 1, 99.8, axis=(0, 1))
-    elif (norm == 'mean-std'):
+    elif norm == "mean-std":
         image = normalize_mean_std(image)
-    elif (norm == 'absolute'):
-        if data_type == '8-bit':
+    elif norm == "absolute":
+        if data_type == "8-bit":
             if normalization_factor is None:
                 image /= 255
             else:
                 image /= normalization_factor
-        elif data_type == '16-bit':
+        elif data_type == "16-bit":
             if normalization_factor is None:
                 image /= 65535
             else:
@@ -611,87 +781,132 @@ def process_one_hot(im, inst, crops_dir, data_subset, crop_size, center, one_hot
     # loop over instances
     ids_subset = np.random.choice(ids, int(fraction_max_ids * len(ids)), replace=False)
     for j, id in enumerate(ids_subset):
-
         y, x = np.where(instance_np[id] == 1)
         ym, xm = np.mean(y), np.mean(x)
 
         jj = int(np.clip(ym - crop_size / 2, 0, h - crop_size))
         ii = int(np.clip(xm - crop_size / 2, 0, w - crop_size))
 
-        if (image[jj:jj + crop_size, ii:ii + crop_size].shape == (crop_size, crop_size) and np.sum(
-                instance_np[id, jj:jj + crop_size, ii:ii + crop_size]) > 0):
-            im_crop = image[jj:jj + crop_size, ii:ii + crop_size]
-            instance_crop = instance_np[:, jj:jj + crop_size, ii:ii + crop_size]
-            center_image_crop = generate_center_image(instance_crop, center, ids, one_hot)
+        if (
+            image[jj : jj + crop_size, ii : ii + crop_size].shape
+            == (crop_size, crop_size)
+            and np.sum(instance_np[id, jj : jj + crop_size, ii : ii + crop_size]) > 0
+        ):
+            im_crop = image[jj : jj + crop_size, ii : ii + crop_size]
+            instance_crop = instance_np[:, jj : jj + crop_size, ii : ii + crop_size]
+            center_image_crop = generate_center_image(
+                instance_crop, center, ids, one_hot
+            )
             instance_crop_sparse = sparsify(instance_crop)
-            tifffile.imsave(image_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j), im_crop)
+            tifffile.imsave(
+                image_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
+                im_crop,
+            )
             if rle_encode:
-                encode(instance_path + os.path.basename(im)[:-4] + "_{:03d}.csv".format(j), instance_crop_sparse,
-                       one_hot=True)
-                encode(center_image_path + os.path.basename(im)[:-4] + "_{:03d}.csv".format(j), center_image_crop,
-                       one_hot=False)
+                encode(
+                    instance_path + os.path.basename(im)[:-4] + "_{:03d}.csv".format(j),
+                    instance_crop_sparse,
+                    one_hot=True,
+                )
+                encode(
+                    center_image_path
+                    + os.path.basename(im)[:-4]
+                    + "_{:03d}.csv".format(j),
+                    center_image_crop,
+                    one_hot=False,
+                )
             else:
-                tifffile.imsave(instance_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
-                                instance_crop_sparse)
-                tifffile.imsave(center_image_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
-                                center_image_crop)
-
+                tifffile.imsave(
+                    instance_path + os.path.basename(im)[:-4] + "_{:03d}.tif".format(j),
+                    instance_crop_sparse,
+                )
+                tifffile.imsave(
+                    center_image_path
+                    + os.path.basename(im)[:-4]
+                    + "_{:03d}.tif".format(j),
+                    center_image_crop,
+                )
+
+
+def process_3d_sliced(
+    im,
+    inst,
+    crops_dir,
+    data_subset,
+    crop_size_x,
+    crop_size_y,
+    crop_size_z,
+    center,
+    one_hot=False,
+    anisotropy_factor=1.0,
+    norm="min-max-percentile",
+    data_type="8-bit",
+    normalization_factor=None,
+    fraction_max_ids=0.01,
+    rle_encode=False,
+    background_id=0,
+):
+    """
+    Processes the actual images and corresponding GT label masks
+    to generate crops of size `crop-size`.
+    Here each volumetric (3D) image is processed sliced-wise
+    to generate 2D crops in XY, YZ and ZX dimensions.
 
-def process_3d_sliced(im, inst, crops_dir, data_subset, crop_size_x, crop_size_y, crop_size_z, center,
-                      one_hot=False, anisotropy_factor=1.0, norm='min-max-percentile', data_type='8-bit',
-                      normalization_factor=None, fraction_max_ids=0.01, rle_encode=False, background_id=0):
+    Parameters
+    ----------
+    im: numpy array
+        Raw image which must be processed (segmented)
+    inst: numpy array
+        Corresponding instance mask which contains objects
+        present in a one-hot encoded style (DYX)
+    crops_dir: string
+        Indicates the path where the crops should be saved
+    crop_size_x: int
+        Length of the crop along the x dimension
+    crop_size_y: int
+        Length of the crop along the y dimension
+    crop_size_z: int
+        Length of the crop along the z dimension
+    center: string,
+        One of `centroid`, `approximate-medoid` or `medoid`
+    one_hot: boolean
+        If True,  each object is encoded as one in it own individual slice
+        in the `inst` image and zero elsewhere.
+        False, by default
+    anisotropy_factor: float
+            This parameter should be set equal to the ratio
+            of the z pixel size to the x or y pixel size
+            Here, we assume that the x or y pixel size is the same
+            (If the imaging is downsampled in the z dimension,
+            `anisotropy_factor` is greater than 1.0)
+    norm: str
+        This should be set equal to one of 'min-max-percentile',
+        'mean-std' or 'absolute'
+    data_type: str
+        This should be set equal to one of '8-bit' or '16-bit'
+        This only comes into play if norm is set equal to 'absolute'
+    normalization_factor: int, None by default
+        This can be set equal to any arbitrary value
+        If not set, then for '8-bit' images, it is taken as 255
+        and for '16-bit' images, it is taken as 65535
+        This only comes into play if norm is set equal to 'absolute'
+    fraction_max_ids: float, between 0 and 1.0
+        This can be set to a small value in case one doesn't wish
+        to produce 2D object centred crops from all objects
+        but only a fraction of objects
+    rle_encode: bool
+        If true, then label masks for crops are not saved as tiffs
+        but as csv files (in a run-length encoded fashion)
+    background_id: int
+        The label id which corresponds to the background
+        This is sset equal to 0 by default
+    Returns
+        -------
     """
-            Processes the actual images and corresponding GT label masks to generate crops of size `crop-size`.
-            Here each volumetric (3D) image is processed sliced-wise to generate 2D crops in XY, YZ and ZX dimensions.
-
-            Parameters
-            ----------
-            im: numpy array
-                Raw image which must be processed (segmented)
-            inst: numpy array
-                Corresponding instance mask which contains objects present in a one-hot encoded style (DYX)
-            crops_dir: string
-                Indicates the path where the crops should be saved
-            crop_size_x: int
-                Length of the crop along the x dimension
-            crop_size_y: int
-                Length of the crop along the y dimension
-            crop_size_z: int
-                Length of the crop along the z dimension
-            center: string,
-                One of `centroid`, `approximate-medoid` or `medoid`
-            one_hot: boolean
-                If True,  each object is encoded as one in it own individual slice in the `inst` image and zero elsewhere.
-                False, by default
-            anisotropy_factor: float
-                    This parameter should be set equal to the ratio of the z pixel size to the x or y pixel size
-                    Here, we assume that the x or y pixel size is the same
-                    (If the imaging is downsampled in the z dimension, `anisotropy_factor` is greater than 1.0)
-            norm: str
-                This should be set equal to one of 'min-max-percentile', 'mean-std' or 'absolute'
-            data_type: str
-                This should be set equal to one of '8-bit' or '16-bit'
-                This only comes into play if norm is set equal to 'absolute'
-            normalization_factor: int, None by default
-                This can be set equal to any arbitrary value
-                If not set, then for '8-bit' images, it is taken as 255
-                and for '16-bit' images, it is taken as 65535
-                This only comes into play if norm is set equal to 'absolute'
-            fraction_max_ids: float, between 0 and 1.0
-                This can be set to a small value in case one doesn't wish to produce 2D object centred crops from all objects
-                but only a fraction of objects
-            rle_encode: bool
-                If true, then label masks for crops are not saved as tiffs but as csv files (in a run-length encoded fashion)
-            background_id: int
-                The label id which corresponds to the background
-                This is sset equal to 0 by default
-            Returns
-                -------
-        """
-
-    image_path = os.path.join(crops_dir, data_subset, 'images/')
-    instance_path = os.path.join(crops_dir, data_subset, 'masks/')
-    center_image_path = os.path.join(crops_dir, data_subset, 'center-' + center + '/')
+
+    image_path = os.path.join(crops_dir, data_subset, "images/")
+    instance_path = os.path.join(crops_dir, data_subset, "masks/")
+    center_image_path = os.path.join(crops_dir, data_subset, "center-" + center + "/")
 
     if not os.path.exists(image_path):
         os.makedirs(os.path.dirname(image_path))
@@ -704,18 +919,18 @@ def process_3d_sliced(im, inst, crops_dir, data_subset, crop_size_x, crop_size_y
         print("Created new directory : {}".format(center_image_path))
 
     instance = tifffile.imread(inst).astype(np.uint16)
-    image = tifffile.imread(im).astype(np.float32)
-    if (norm == 'min-max-percentile'):
+    image = tifffile.imread(im).astype(float)
+    if norm == "min-max-percentile":
         image = normalize_min_max_percentile(image, 1, 99.8, axis=(0, 1, 2))
-    elif (norm == 'mean-std'):
+    elif norm == "mean-std":
         image = normalize_mean_std(image)
-    elif (norm == 'absolute'):
-        if data_type == '8-bit':
+    elif norm == "absolute":
+        if data_type == "8-bit":
             if normalization_factor is None:
                 image /= 255
             else:
                 image /= normalization_factor
-        elif data_type == '16-bit':
+        elif data_type == "16-bit":
             if normalization_factor is None:
                 image /= 65535
             else:
@@ -736,7 +951,9 @@ def process_3d_sliced(im, inst, crops_dir, data_subset, crop_size_x, crop_size_y
     ids = ids[ids != 0]
 
     # loop over instances
-    ids_subset = np.random.choice(ids, int(np.ceil(fraction_max_ids * len(ids))), replace=False)
+    ids_subset = np.random.choice(
+        ids, int(np.ceil(fraction_max_ids * len(ids))), replace=False
+    )
     for id in ids_subset:
         z, y, x = np.where(instance_np == id)
         zm, ym, xm = np.mean(z), np.mean(y), np.mean(x)
@@ -746,127 +963,248 @@ def process_3d_sliced(im, inst, crops_dir, data_subset, crop_size_x, crop_size_y
 
         # ZY
         for x in range(ii, ii + crop_size_x, int(np.ceil(anisotropy_factor))):
-            if (image[kk:kk + crop_size_z, jj:jj + crop_size_y, x].shape == (crop_size_z, crop_size_y)):
-                im_crop = image[kk:kk + crop_size_z, jj:jj + crop_size_y, x]
-                instance_crop = instance_np[kk:kk + crop_size_z, jj:jj + crop_size_y, x]
+            if image[kk : kk + crop_size_z, jj : jj + crop_size_y, x].shape == (
+                crop_size_z,
+                crop_size_y,
+            ):
+                im_crop = image[kk : kk + crop_size_z, jj : jj + crop_size_y, x]
+                instance_crop = instance_np[
+                    kk : kk + crop_size_z, jj : jj + crop_size_y, x
+                ]
                 if np.sum(instance_crop) > 0:  # shouldn't be only background
-                    center_image_crop = generate_center_image(instance_crop, center, ids, one_hot=one_hot)
+                    center_image_crop = generate_center_image(
+                        instance_crop, center, ids, one_hot=one_hot
+                    )
                     tifffile.imsave(
-                        image_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                            x) + "_ZY.tif",
-                        im_crop)
+                        image_path
+                        + os.path.basename(im)[:-4]
+                        + "_{:03d}".format(id)
+                        + "_{:03d}".format(x)
+                        + "_ZY.tif",
+                        im_crop,
+                    )
                     if rle_encode:
-                        encode(instance_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                            x) + "_ZY.csv", instance_crop.astype(np.uint16))
-                        encode(center_image_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                            x) + "_ZY.csv", center_image_crop)
+                        encode(
+                            instance_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(x)
+                            + "_ZY.csv",
+                            instance_crop.astype(np.uint16),
+                        )
+                        encode(
+                            center_image_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(x)
+                            + "_ZY.csv",
+                            center_image_crop,
+                        )
                     else:
                         tifffile.imsave(
-                            instance_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                                x) + "_ZY.tif", instance_crop.astype(np.uint16))
+                            instance_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(x)
+                            + "_ZY.tif",
+                            instance_crop.astype(np.uint16),
+                        )
                         tifffile.imsave(
-                            center_image_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                                x) + "_ZY.tif", center_image_crop)
+                            center_image_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(x)
+                            + "_ZY.tif",
+                            center_image_crop,
+                        )
 
         # YX
         for z in range(kk, kk + crop_size_z, int(np.ceil(anisotropy_factor))):
-            if (image[z, jj:jj + crop_size_y, ii:ii + crop_size_x].shape == (crop_size_y, crop_size_x)):
-                im_crop = image[z, jj:jj + crop_size_y, ii:ii + crop_size_x]
-                instance_crop = instance_np[z, jj:jj + crop_size_y, ii:ii + crop_size_x]
+            if image[z, jj : jj + crop_size_y, ii : ii + crop_size_x].shape == (
+                crop_size_y,
+                crop_size_x,
+            ):
+                im_crop = image[z, jj : jj + crop_size_y, ii : ii + crop_size_x]
+                instance_crop = instance_np[
+                    z, jj : jj + crop_size_y, ii : ii + crop_size_x
+                ]
                 if np.sum(instance_crop) > 0:
-                    center_image_crop = generate_center_image(instance_crop, center, ids, one_hot=one_hot)
+                    center_image_crop = generate_center_image(
+                        instance_crop, center, ids, one_hot=one_hot
+                    )
                     tifffile.imsave(
-                        image_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                            z) + "_YX.tif",
-                        im_crop)
+                        image_path
+                        + os.path.basename(im)[:-4]
+                        + "_{:03d}".format(id)
+                        + "_{:03d}".format(z)
+                        + "_YX.tif",
+                        im_crop,
+                    )
                     if rle_encode:
-                        encode(instance_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                            z) + "_YX.csv", instance_crop.astype(np.uint16))
-                        encode(center_image_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                            z) + "_YX.csv", center_image_crop)
+                        encode(
+                            instance_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(z)
+                            + "_YX.csv",
+                            instance_crop.astype(np.uint16),
+                        )
+                        encode(
+                            center_image_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(z)
+                            + "_YX.csv",
+                            center_image_crop,
+                        )
                     else:
                         tifffile.imsave(
-                            instance_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                                z) + "_YX.tif", instance_crop.astype(np.uint16))
+                            instance_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(z)
+                            + "_YX.tif",
+                            instance_crop.astype(np.uint16),
+                        )
                         tifffile.imsave(
-                            center_image_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                                z) + "_YX.tif", center_image_crop)
+                            center_image_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(z)
+                            + "_YX.tif",
+                            center_image_crop,
+                        )
 
         # XZ
         for y in range(jj, jj + crop_size_y, int(np.ceil(anisotropy_factor))):
-            if (image[kk:kk + crop_size_z, y, ii:ii + crop_size_x].shape == (crop_size_z, crop_size_x)):
-                im_crop = np.transpose(image[kk:kk + crop_size_z, y, ii:ii + crop_size_x])
-                instance_crop = np.transpose(instance_np[kk:kk + crop_size_z, y, ii:ii + crop_size_x])
+            if image[kk : kk + crop_size_z, y, ii : ii + crop_size_x].shape == (
+                crop_size_z,
+                crop_size_x,
+            ):
+                im_crop = np.transpose(
+                    image[kk : kk + crop_size_z, y, ii : ii + crop_size_x]
+                )
+                instance_crop = np.transpose(
+                    instance_np[kk : kk + crop_size_z, y, ii : ii + crop_size_x]
+                )
                 if np.sum(instance_crop) > 0:
-                    center_image_crop = generate_center_image(instance_crop, center, ids, one_hot=one_hot)
+                    center_image_crop = generate_center_image(
+                        instance_crop, center, ids, one_hot=one_hot
+                    )
                     tifffile.imsave(
-                        image_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                            y) + "_XZ.tif",
-                        im_crop)
+                        image_path
+                        + os.path.basename(im)[:-4]
+                        + "_{:03d}".format(id)
+                        + "_{:03d}".format(y)
+                        + "_XZ.tif",
+                        im_crop,
+                    )
                     if rle_encode:
-                        encode(instance_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                            y) + "_XZ.csv", instance_crop.astype(np.uint16))
-                        encode(center_image_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                            y) + "_XZ.csv", center_image_crop)
+                        encode(
+                            instance_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(y)
+                            + "_XZ.csv",
+                            instance_crop.astype(np.uint16),
+                        )
+                        encode(
+                            center_image_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(y)
+                            + "_XZ.csv",
+                            center_image_crop,
+                        )
                     else:
                         tifffile.imsave(
-                            instance_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                                y) + "_XZ.tif", instance_crop.astype(np.uint16))
+                            instance_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(y)
+                            + "_XZ.tif",
+                            instance_crop.astype(np.uint16),
+                        )
                         tifffile.imsave(
-                            center_image_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                                y) + "_XZ.tif", center_image_crop)
-
+                            center_image_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(y)
+                            + "_XZ.tif",
+                            center_image_crop,
+                        )
+
+
+def process_3d_ilp(
+    im,
+    inst,
+    crops_dir,
+    data_subset,
+    crop_size_x,
+    crop_size_y,
+    center,
+    one_hot=False,
+    norm="min-max-percentile",
+    data_type="8-bit",
+    normalization_factor=None,
+    fraction_max_ids=0.10,
+    rle_encode=False,
+    background_id=0,
+):
+    """
+    Processes the actual images and corresponding GT label masks
+    to generate crops of size `crop-size`.
+    Here each volumetric (3D) image is processed sliced-wise
+    to generate 2D crops in XY dimensions.
 
-def process_3d_ilp(im, inst, crops_dir, data_subset, crop_size_x, crop_size_y, center,
-                   one_hot=False, norm='min-max-percentile', data_type='8-bit',
-                   normalization_factor=None, fraction_max_ids=0.10, rle_encode=False, background_id=0):
+    Parameters
+    ----------
+    im: numpy array
+        Raw image which must be processed (segmented)
+    inst: numpy array
+        Corresponding instance mask which contains objects present
+        in a one-hot encoded style (DYX)
+    crops_dir: string
+        Indicates the path where the crops should be saved
+    crop_size_y: int
+        Height of the crop in the y dimension
+    crop_size_x: int
+        Width of the crop in the x dimension
+    center: string,
+        One of `centroid`, `approximate-medoid` or `medoid`
+    norm: boolean
+        Setting this to True would perfom min-max normalization
+        on the raw image prior to cropping them
+        False by default
+    one_hot: boolean
+        If True,  each object is encoded as one in it own individual slice
+        in the `inst` image and zero elsewhere.
+        False, by default
+    data_type: str
+        This should be set equal to one of '8-bit' or '16-bit'
+        This only comes into play if norm is set equal to 'absolute'
+    normalization_factor: int, None by default
+        This can be set equal to any arbitrary value
+        If not set, then for '8-bit' images, it is taken as 255
+        and for '16-bit' images, it is taken as 65535
+        This only comes into play if norm is set equal to 'absolute'
+    fraction_max_ids: float, between 0 and 1.0
+        This can be set to a small value in case one doesn't wish to produce
+        2D object centred crops from all objects
+        but only a fraction of objects
+    rle_encode: bool
+        If true, then label masks for crops are not saved as tiffs
+        but as csv files (in a run-length encoded fashion)
+    background_id: int
+        The label id which corresponds to the background
+        This is sset equal to 0 by default
+    Returns
+        -------
     """
-            Processes the actual images and corresponding GT label masks to generate crops of size `crop-size`.
-            Here each volumetric (3D) image is processed sliced-wise to generate 2D crops in XY dimensions.
-
-            Parameters
-            ----------
-            im: numpy array
-                Raw image which must be processed (segmented)
-            inst: numpy array
-                Corresponding instance mask which contains objects present in a one-hot encoded style (DYX)
-            crops_dir: string
-                Indicates the path where the crops should be saved
-            crop_size_y: int
-                Height of the crop in the y dimension
-            crop_size_x: int
-                Width of the crop in the x dimension
-            center: string,
-                One of `centroid`, `approximate-medoid` or `medoid`
-            norm: boolean
-                Setting this to True would perfom min-max normalization on the raw image prior to cropping them
-                False by default
-            one_hot: boolean
-                If True,  each object is encoded as one in it own individual slice in the `inst` image and zero elsewhere.
-                False, by default
-            data_type: str
-                This should be set equal to one of '8-bit' or '16-bit'
-                This only comes into play if norm is set equal to 'absolute'
-            normalization_factor: int, None by default
-                This can be set equal to any arbitrary value
-                If not set, then for '8-bit' images, it is taken as 255
-                and for '16-bit' images, it is taken as 65535
-                This only comes into play if norm is set equal to 'absolute'
-            fraction_max_ids: float, between 0 and 1.0
-                This can be set to a small value in case one doesn't wish to produce 2D object centred crops from all objects
-                but only a fraction of objects
-            rle_encode: bool
-                If true, then label masks for crops are not saved as tiffs but as csv files (in a run-length encoded fashion)
-            background_id: int
-                The label id which corresponds to the background
-                This is sset equal to 0 by default
-            Returns
-                -------
-        """
-
-    image_path = os.path.join(crops_dir, data_subset, 'images/')
-    instance_path = os.path.join(crops_dir, data_subset, 'masks/')
-    center_image_path = os.path.join(crops_dir, data_subset, 'center-' + center + '/')
+
+    image_path = os.path.join(crops_dir, data_subset, "images/")
+    instance_path = os.path.join(crops_dir, data_subset, "masks/")
+    center_image_path = os.path.join(crops_dir, data_subset, "center-" + center + "/")
 
     if not os.path.exists(image_path):
         os.makedirs(os.path.dirname(image_path))
@@ -879,18 +1217,18 @@ def process_3d_ilp(im, inst, crops_dir, data_subset, crop_size_x, crop_size_y, c
         print("Created new directory : {}".format(center_image_path))
 
     instance = tifffile.imread(inst).astype(np.uint16)
-    image = tifffile.imread(im).astype(np.float32)
-    if (norm == 'min-max-percentile'):
+    image = tifffile.imread(im).astype(float)
+    if norm == "min-max-percentile":
         image = normalize_min_max_percentile(image, 1, 99.8, axis=(0, 1, 2))
-    elif (norm == 'mean-std'):
+    elif norm == "mean-std":
         image = normalize_mean_std(image)
-    elif (norm == 'absolute'):
-        if data_type == '8-bit':
+    elif norm == "absolute":
+        if data_type == "8-bit":
             if normalization_factor is None:
                 image /= 255
             else:
                 image /= normalization_factor
-        elif data_type == '16-bit':
+        elif data_type == "16-bit":
             if normalization_factor is None:
                 image /= 65535
             else:
@@ -906,30 +1244,65 @@ def process_3d_ilp(im, inst, crops_dir, data_subset, crop_size_x, crop_size_y, c
         instance_np[z][instance_np[z] == background_id] = 0
         ids = np.unique(instance_np[z][object_mask])
         ids = ids[ids != 0]
-        ids_subset = np.random.choice(ids, int(fraction_max_ids * len(ids)), replace=False)
+        ids_subset = np.random.choice(
+            ids, int(fraction_max_ids * len(ids)), replace=False
+        )
         for id in ids_subset:
             y, x = np.where(instance_np[z] == id)
             ym, xm = np.mean(y), np.mean(x)
             jj = int(np.clip(ym - crop_size_y / 2, 0, h - crop_size_y))
             ii = int(np.clip(xm - crop_size_x / 2, 0, w - crop_size_x))
-            if (image[z, jj:jj + crop_size_y, ii:ii + crop_size_x].shape == (crop_size_y, crop_size_x)):
-                im_crop = image[z, jj:jj + crop_size_y, ii:ii + crop_size_x]
-                instance_crop = instance_np[z, jj:jj + crop_size_y, ii:ii + crop_size_x]
+            if image[z, jj : jj + crop_size_y, ii : ii + crop_size_x].shape == (
+                crop_size_y,
+                crop_size_x,
+            ):
+                im_crop = image[z, jj : jj + crop_size_y, ii : ii + crop_size_x]
+                instance_crop = instance_np[
+                    z, jj : jj + crop_size_y, ii : ii + crop_size_x
+                ]
                 if np.sum(instance_crop) > 0:
-                    center_image_crop = generate_center_image(instance_crop, center, ids, one_hot=one_hot)
+                    center_image_crop = generate_center_image(
+                        instance_crop, center, ids, one_hot=one_hot
+                    )
                     tifffile.imsave(
-                        image_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                            z) + "_YX.tif",
-                        im_crop)
+                        image_path
+                        + os.path.basename(im)[:-4]
+                        + "_{:03d}".format(id)
+                        + "_{:03d}".format(z)
+                        + "_YX.tif",
+                        im_crop,
+                    )
                     if rle_encode:
-                        encode(instance_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                            z) + "_YX.csv", instance_crop.astype(np.uint16))
-                        encode(center_image_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                            z) + "_YX.csv", center_image_crop)
+                        encode(
+                            instance_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(z)
+                            + "_YX.csv",
+                            instance_crop.astype(np.uint16),
+                        )
+                        encode(
+                            center_image_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(z)
+                            + "_YX.csv",
+                            center_image_crop,
+                        )
                     else:
                         tifffile.imsave(
-                            instance_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                                z) + "_YX.tif", instance_crop.astype(np.uint16))
+                            instance_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(z)
+                            + "_YX.tif",
+                            instance_crop.astype(np.uint16),
+                        )
                         tifffile.imsave(
-                            center_image_path + os.path.basename(im)[:-4] + "_{:03d}".format(id) + "_{:03d}".format(
-                                z) + "_YX.tif", center_image_crop)
+                            center_image_path
+                            + os.path.basename(im)[:-4]
+                            + "_{:03d}".format(id)
+                            + "_{:03d}".format(z)
+                            + "_YX.tif",
+                            center_image_crop,
+                        )
diff --git a/EmbedSeg/utils/glasbey.py b/EmbedSeg/utils/glasbey.py
index 5a90768..464de88 100644
--- a/EmbedSeg/utils/glasbey.py
+++ b/EmbedSeg/utils/glasbey.py
@@ -30,15 +30,12 @@
     class Bar:
         pass
 
-
     class ETA:
         pass
 
-
     class Percentage:
         pass
 
-
     class ProgressBar:
         def __init__(self, **kwargs):
             pass
@@ -54,13 +51,15 @@ def finish(self):
 
 
 class Glasbey:
-    def __init__(self,
-                 base_palette=None,
-                 overwrite_base_palette: bool = False,
-                 no_black: bool = False,
-                 lightness_range=None,
-                 chroma_range=None,
-                 hue_range=None):
+    def __init__(
+        self,
+        base_palette=None,
+        overwrite_base_palette: bool = False,
+        no_black: bool = False,
+        lightness_range=None,
+        chroma_range=None,
+        hue_range=None,
+    ):
         # Constants
         self.MAX = 256
         self.NUM_COLORS = self.MAX * self.MAX * self.MAX
@@ -70,12 +69,20 @@ def __init__(self,
 
         # Check input
         if type(base_palette) == str:
-            assert os.path.isfile(base_palette), "file does not exist: {}".format(base_palette)
+            assert os.path.isfile(base_palette), "file does not exist: {}".format(
+                base_palette
+            )
         elif type(base_palette) == list:
-            assert self.check_validity_rbg_palette(base_palette), "Base palette must be in this format: [(255,255,255), ...]"
-            assert not self.overwrite_base_palette, "base_palette is no file, cannot overwrite it!"
+            assert self.check_validity_rbg_palette(
+                base_palette
+            ), "Base palette must be in this format: [(255,255,255), ...]"
+            assert (
+                not self.overwrite_base_palette
+            ), "base_palette is no file, cannot overwrite it!"
         else:
-            assert not self.overwrite_base_palette, "no base_palette specified, cannot overwrite it!"
+            assert (
+                not self.overwrite_base_palette
+            ), "no base_palette specified, cannot overwrite it!"
 
         # Load colors
         self.colors = self.load_or_generate_color_table()
@@ -86,35 +93,39 @@ def __init__(self,
             self.palette = self.load_palette(base_palette)
             self.palette = [self.colors[i, :] for i in self.palette]
         elif type(base_palette) == list and len(base_palette) > 0:
-            self.palette = [(rgb[0] * 256 + rgb[1]) * 256 + rgb[2] for rgb in base_palette]
+            self.palette = [
+                (rgb[0] * 256 + rgb[1]) * 256 + rgb[2] for rgb in base_palette
+            ]
             self.palette = [self.colors[i, :] for i in self.palette]
         else:
             self.palette = [self.colors[-1, :]]  # white
 
-        assert self.check_validity_internal_palette(), "Internal error during __init__: self.palette is poorly formatted."
+        assert (
+            self.check_validity_internal_palette()
+        ), "Internal error during __init__: self.palette is poorly formatted."
 
         # Update self.colors
         # Exclude greys (values with low Chroma in JCh) and set lightness range,
         if lightness_range is not None:
             jch = cspace_convert(self.colors, "CAM02-UCS", "JCh")
             self.colors = self.colors[
-                          (jch[:, 0] >= lightness_range[0]) & (jch[:, 0] <= lightness_range[1]), :
-                          ]
+                (jch[:, 0] >= lightness_range[0]) & (jch[:, 0] <= lightness_range[1]), :
+            ]
         if chroma_range is not None:
             jch = cspace_convert(self.colors, "CAM02-UCS", "JCh")
             self.colors = self.colors[
-                          (jch[:, 1] >= chroma_range[0]) & (jch[:, 1] <= chroma_range[1]), :
-                          ]
+                (jch[:, 1] >= chroma_range[0]) & (jch[:, 1] <= chroma_range[1]), :
+            ]
         if hue_range is not None:
             jch = cspace_convert(self.colors, "CAM02-UCS", "JCh")
             if hue_range[0] > hue_range[1]:
                 self.colors = self.colors[
-                              (jch[:, 2] >= hue_range[0]) | (jch[:, 2] <= hue_range[1]), :
-                              ]
+                    (jch[:, 2] >= hue_range[0]) | (jch[:, 2] <= hue_range[1]), :
+                ]
             else:
                 self.colors = self.colors[
-                              (jch[:, 2] >= hue_range[0]) & (jch[:, 2] <= hue_range[1]), :
-                              ]
+                    (jch[:, 2] >= hue_range[0]) & (jch[:, 2] <= hue_range[1]), :
+                ]
         # Exclude colors that are close to black
         if no_black:
             MIN_DISTANCE_TO_BLACK = 35
@@ -153,10 +164,17 @@ def update_distances(colors, color):
             pbar.update(len(self.palette))
         pbar.finish()
 
-        assert self.check_validity_internal_palette(), "Internal error during extend_palette: self.palette is poorly formatted."
+        assert (
+            self.check_validity_internal_palette()
+        ), "Internal error during extend_palette: self.palette is poorly formatted."
 
         if self.overwrite_base_palette:
-            self.save_palette(palette=self.palette, path=self.base_palette, format="byte", overwrite=True)
+            self.save_palette(
+                palette=self.palette,
+                path=self.base_palette,
+                format="byte",
+                overwrite=True,
+            )
 
         return cspace_convert(self.palette[0:size], "CAM02-UCS", "sRGB1")
 
@@ -190,8 +208,8 @@ def generate_color_table(self):
                 d = i * self.MAX
                 for b in range(self.MAX):
                     colors[d + b, :] = (r, g, b)
-                colors[d:d + self.MAX] = cspace_convert(
-                    colors[d:d + self.MAX], "sRGB255", "CAM02-UCS"
+                colors[d : d + self.MAX] = cspace_convert(
+                    colors[d : d + self.MAX], "sRGB255", "CAM02-UCS"
                 )
                 pbar.update(i)
                 i += 1
@@ -205,7 +223,7 @@ def load_palette(path):
         """
         assert os.path.isfile(path)
         palette = list()
-        with open(path, 'r') as file:
+        with open(path, "r") as file:
             for line in file:
                 rgb = [int(c) for c in line.strip().split(",")]
                 palette.append((rgb[0] * 256 + rgb[1]) * 256 + rgb[2])
@@ -221,20 +239,24 @@ def save_palette(palette, path: str, format: str = "byte", overwrite: bool = Fal
         if not overwrite:
             assert not os.path.isfile(path)
 
-        with open(path, 'w') as file:
+        with open(path, "w") as file:
             if format.lower() == "byte":
                 for color in palette:
                     rgb255 = tuple(int(round(k * 255)) for k in color)
                     file.write("{},{},{}\n".format(*rgb255))
             elif format.lower() == "float":
                 for color in palette:
-                    file.write("{:.6f},{:.6f},{:.6f}\n".format(*(abs(k) for k in color)))
+                    file.write(
+                        "{:.6f},{:.6f},{:.6f}\n".format(*(abs(k) for k in color))
+                    )
             elif format.lower() == "hex":
                 for color in palette:
                     rgb255 = tuple(int(round(k * 255)) for k in color)
                     file.write("#{:02x}{:02x}{:02x}\n".format(*rgb255))
             else:
-                raise ValueError("Format doesn't match. Choose between 'byte', 'hex', and 'float'")
+                raise ValueError(
+                    "Format doesn't match. Choose between 'byte', 'hex', and 'float'"
+                )
 
     def check_validity_internal_palette(self):
         if type(self.palette) != list:
@@ -251,7 +273,11 @@ def check_validity_rbg_palette(palette):
         for color in palette:
             if len(color) != 3 or type(color) != tuple:
                 return False
-            if not 0 <= color[0] <= 255 and 0 <= color[1] <= 255 and 0 <= color[2] <= 255:
+            if (
+                not 0 <= color[0] <= 255
+                and 0 <= color[1] <= 255
+                and 0 <= color[2] <= 255
+            ):
                 return False
 
         return True
@@ -327,27 +353,35 @@ def view_palette(palette):
         type=ast.literal_eval,
         help="set min and max for chroma (e.g. 10,100)",
     )
-    parser.add_argument("--hue-range", type=ast.literal_eval,
-                        help="set start and end for hue (e.g. 315,45)")
-    parser.add_argument("--view", action="store_true",
-                        help="view generated palette")
-    parser.add_argument("--format", default="byte", choices=["byte", "hex", "float"],
-                        help="output format")
-    parser.add_argument("size", type=int,
-                        help="number of colors in the palette")
-    parser.add_argument("output", type=argparse.FileType("w"),
-                        help="output palette filename")
+    parser.add_argument(
+        "--hue-range",
+        type=ast.literal_eval,
+        help="set start and end for hue (e.g. 315,45)",
+    )
+    parser.add_argument("--view", action="store_true", help="view generated palette")
+    parser.add_argument(
+        "--format",
+        default="byte",
+        choices=["byte", "hex", "float"],
+        help="output format",
+    )
+    parser.add_argument("size", type=int, help="number of colors in the palette")
+    parser.add_argument(
+        "output", type=argparse.FileType("w"), help="output palette filename"
+    )
     args = parser.parse_args()
 
     if args.format not in ["byte", "hex", "float"]:
         sys.exit('Invalid output format "{}"'.format(args.format))
 
-    gb = Glasbey(base_palette=getattr(args.base_palette, "name", None),
-                 overwrite_base_palette=False,
-                 no_black=args.no_black,
-                 lightness_range=args.lightness_range,
-                 chroma_range=args.chroma_range,
-                 hue_range=args.hue_range)
+    gb = Glasbey(
+        base_palette=getattr(args.base_palette, "name", None),
+        overwrite_base_palette=False,
+        no_black=args.no_black,
+        lightness_range=args.lightness_range,
+        chroma_range=args.chroma_range,
+        hue_range=args.hue_range,
+    )
 
     new_palette = gb.generate_palette(size=args.size)
     assert len(new_palette) == args.size
diff --git a/EmbedSeg/utils/metrics.py b/EmbedSeg/utils/metrics.py
index 9a34927..31c17a1 100644
--- a/EmbedSeg/utils/metrics.py
+++ b/EmbedSeg/utils/metrics.py
@@ -14,16 +14,19 @@
 
 
 def label_are_sequential(y):
-    """ returns true if y has only sequential labels from 1... """
+    """returns true if y has only sequential labels from 1..."""
     labels = np.unique(y)
     return (set(labels) - {0}) == set(range(1, 1 + labels.max()))
 
 
 def _check_label_array(y, name=None, check_sequential=False):
-    err = ValueError("{label} must be an array of {integers}.".format(
-        label='labels' if name is None else name,
-        integers=('sequential ' if check_sequential else '') + 'non-negative integers',
-    ))
+    err = ValueError(
+        "{label} must be an array of {integers}.".format(
+            label="labels" if name is None else name,
+            integers=("sequential " if check_sequential else "")
+            + "non-negative integers",
+        )
+    )
     is_array_of_integers(y) or _raise(err)
     if check_sequential:
         label_are_sequential(y) or _raise(err)
@@ -80,15 +83,15 @@ def relabel_sequential(label_field, offset=1):
     if not (labels == 0).any():
         labels = np.concatenate(([0], labels))
     inverse_map = np.zeros(offset - 1 + len(labels), dtype=label_field.dtype)
-    inverse_map[(offset - 1):] = labels
+    inverse_map[(offset - 1) :] = labels
     relabeled = forward_map[label_field]
     return relabeled, forward_map, inverse_map
 
 
 def label_overlap(x, y, check=True):
     if check:
-        _check_label_array(x, 'x', True)
-        _check_label_array(y, 'y', True)
+        _check_label_array(x, "x", True)
+        _check_label_array(y, "y", True)
         x.shape == y.shape or _raise(ValueError("x and y must have the same shape"))
     return _label_overlap(x, y)
 
@@ -104,7 +107,7 @@ def _label_overlap(x, y):
 
 
 def intersection_over_union(overlap):
-    _check_label_array(overlap, 'overlap')
+    _check_label_array(overlap, "overlap")
     if np.sum(overlap) == 0:
         return overlap
     n_pixels_pred = np.sum(overlap, axis=0, keepdims=True)
@@ -112,20 +115,55 @@ def intersection_over_union(overlap):
     return overlap / (n_pixels_pred + n_pixels_true - overlap)
 
 
-matching_criteria['iou'] = intersection_over_union
+matching_criteria["iou"] = intersection_over_union
 
 
-def matching_dataset(y_true, y_pred, thresh=0.5, criterion='iou', by_image=False, show_progress=True, parallel=False):
-    len(y_true) == len(y_pred) or _raise(ValueError("y_true and y_pred must have the same length."))
+def matching_dataset(
+    y_true,
+    y_pred,
+    thresh=0.5,
+    criterion="iou",
+    by_image=False,
+    show_progress=True,
+    parallel=False,
+):
+    len(y_true) == len(y_pred) or _raise(
+        ValueError("y_true and y_pred must have the same length.")
+    )
     return matching_dataset_lazy(
-        tuple(zip(y_true, y_pred)), thresh=thresh, criterion=criterion, by_image=by_image, show_progress=show_progress,
+        tuple(zip(y_true, y_pred)),
+        thresh=thresh,
+        criterion=criterion,
+        by_image=by_image,
+        show_progress=show_progress,
         parallel=parallel,
     )
 
 
-def matching_dataset_lazy(y_gen, thresh=0.5, criterion='iou', by_image=False, show_progress=True, parallel=False):
-    expected_keys = set(('fp', 'tp', 'fn', 'precision', 'recall', 'accuracy', 'f1', 'criterion', 'thresh', 'n_true',
-                         'n_pred', 'mean_true_score'))
+def matching_dataset_lazy(
+    y_gen,
+    thresh=0.5,
+    criterion="iou",
+    by_image=False,
+    show_progress=True,
+    parallel=False,
+):
+    expected_keys = set(
+        (
+            "fp",
+            "tp",
+            "fn",
+            "precision",
+            "recall",
+            "accuracy",
+            "f1",
+            "criterion",
+            "thresh",
+            "n_true",
+            "n_pred",
+            "mean_true_score",
+        )
+    )
 
     single_thresh = False
     if np.isscalar(thresh):
@@ -133,14 +171,19 @@ def matching_dataset_lazy(y_gen, thresh=0.5, criterion='iou', by_image=False, sh
         thresh = (thresh,)
 
     tqdm_kwargs = {}
-    tqdm_kwargs['disable'] = not bool(show_progress)
+    tqdm_kwargs["disable"] = not bool(show_progress)
     if int(show_progress) > 1:
-        tqdm_kwargs['total'] = int(show_progress)
+        tqdm_kwargs["total"] = int(show_progress)
 
     # compute matching stats for every pair of label images
     if parallel:
         from concurrent.futures import ThreadPoolExecutor
-        fn = lambda pair: matching(*pair, thresh=thresh, criterion=criterion, report_matches=False)
+
+        def fn(pair):
+            return matching(
+                *pair, thresh=thresh, criterion=criterion, report_matches=False
+            )
+
         with ThreadPoolExecutor() as pool:
             stats_all = tuple(pool.map(fn, tqdm(y_gen, **tqdm_kwargs)))
     else:
@@ -156,7 +199,7 @@ def matching_dataset_lazy(y_gen, thresh=0.5, criterion='iou', by_image=False, sh
         for i, s in enumerate(stats):
             acc = accumulate[i]
             for k, v in s._asdict().items():
-                if k == 'mean_true_score' and not bool(by_image):
+                if k == "mean_true_score" and not bool(by_image):
                     # convert mean_true_score to "sum_true_score"
                     acc[k] = acc.setdefault(k, 0) + v * s.n_true
                 else:
@@ -168,38 +211,51 @@ def matching_dataset_lazy(y_gen, thresh=0.5, criterion='iou', by_image=False, sh
     # normalize/compute 'precision', 'recall', 'accuracy', 'f1'
     for thr, acc in zip(thresh, accumulate):
         set(acc.keys()) == expected_keys or _raise(ValueError("unexpected keys"))
-        acc['criterion'] = criterion
-        acc['thresh'] = thr
-        acc['by_image'] = bool(by_image)
+        acc["criterion"] = criterion
+        acc["thresh"] = thr
+        acc["by_image"] = bool(by_image)
         if bool(by_image):
-            for k in ('precision', 'recall', 'accuracy', 'f1', 'mean_true_score'):
+            for k in ("precision", "recall", "accuracy", "f1", "mean_true_score"):
                 acc[k] /= n_images
         else:
-            tp, fp, fn = acc['tp'], acc['fp'], acc['fn']
+            tp, fp, fn = acc["tp"], acc["fp"], acc["fn"]
             acc.update(
                 precision=precision(tp, fp, fn),
                 recall=recall(tp, fp, fn),
                 accuracy=accuracy(tp, fp, fn),
                 f1=f1(tp, fp, fn),
-                mean_true_score=acc['mean_true_score'] / acc['n_true'] if acc['n_true'] > 0 else 0.0,
+                mean_true_score=acc["mean_true_score"] / acc["n_true"]
+                if acc["n_true"] > 0
+                else 0.0,
             )
 
-    accumulate = tuple(namedtuple('DatasetMatching', acc.keys())(*acc.values()) for acc in accumulate)
+    accumulate = tuple(
+        namedtuple("DatasetMatching", acc.keys())(*acc.values()) for acc in accumulate
+    )
     return accumulate[0] if single_thresh else accumulate
 
 
-def matching(y_true, y_pred, thresh=0.5, criterion='iou', report_matches=False):
+def matching(y_true, y_pred, thresh=0.5, criterion="iou", report_matches=False):
     """
-    if report_matches=True, return (matched_pairs,matched_scores) are independent of 'thresh'
+    if report_matches=True,
+    return (matched_pairs,matched_scores) are independent of 'thresh'
     """
 
-    _check_label_array(y_true, 'y_true')
-    _check_label_array(y_pred, 'y_pred')
-    y_true.shape == y_pred.shape or _raise(ValueError(
-        "y_true ({y_true.shape}) and y_pred ({y_pred.shape}) have different shapes".format(y_true=y_true,
-                                                                                           y_pred=y_pred)))
-    criterion in matching_criteria or _raise(ValueError("Matching criterion '%s' not supported." % criterion))
-    if thresh is None: thresh = 0
+    _check_label_array(y_true, "y_true")
+    _check_label_array(y_pred, "y_pred")
+    y_true.shape == y_pred.shape or _raise(
+        ValueError(
+            "y_true ({y_true.shape}) and y_pred ({y_pred.shape}) \
+                    have diff shapes".format(
+                y_true=y_true, y_pred=y_pred
+            )
+        )
+    )
+    criterion in matching_criteria or _raise(
+        ValueError("Matching criterion '%s' not supported." % criterion)
+    )
+    if thresh is None:
+        thresh = 0
     thresh = float(thresh) if np.isscalar(thresh) else map(float, thresh)
 
     y_true, _, map_rev_true = relabel_sequential(y_true)
@@ -239,14 +295,18 @@ def _single(thr):
             f1=f1(tp, fp, fn),
             n_true=n_true,
             n_pred=n_pred,
-            mean_true_score=np.sum(scores[true_ind, pred_ind][match_ok]) / n_true if not_trivial else 0.0,
+            mean_true_score=np.sum(scores[true_ind, pred_ind][match_ok]) / n_true
+            if not_trivial
+            else 0.0,
         )
         if bool(report_matches):
             if not_trivial:
                 stats_dict.update(
                     # int() to be json serializable
                     matched_pairs=tuple(
-                        (int(map_rev_true[i]), int(map_rev_pred[j])) for i, j in zip(1 + true_ind, 1 + pred_ind)),
+                        (int(map_rev_true[i]), int(map_rev_pred[j]))
+                        for i, j in zip(1 + true_ind, 1 + pred_ind)
+                    ),
                     matched_scores=tuple(scores[true_ind, pred_ind]),
                     matched_tps=tuple(map(int, np.flatnonzero(match_ok))),
                 )
@@ -256,31 +316,31 @@ def _single(thr):
                     matched_scores=(),
                     matched_tps=(),
                 )
-        return namedtuple('Matching', stats_dict.keys())(*stats_dict.values())
+        return namedtuple("Matching", stats_dict.keys())(*stats_dict.values())
 
     return _single(thresh) if np.isscalar(thresh) else tuple(map(_single, thresh))
 
 
 def obtain_APdsb_one_hot(gt_image, prediction_image, ap_val):
     """
-        Obtain Average Precision (AP_dsb) or Accuracy for GT, one-hot label masks
+    Obtain Average Precision (AP_dsb) or Accuracy for GT, one-hot label masks
 
-        Parameters
-        -------
+    Parameters
+    -------
 
-        gt_image: numpy array
-            label mask (IYX)
-        prediction_image: numpy array
-            label mask (YX)
-        ap_val: float
-            IoU Threshold between 0 and 1.0
+    gt_image: numpy array
+        label mask (IYX)
+    prediction_image: numpy array
+        label mask (YX)
+    ap_val: float
+        IoU Threshold between 0 and 1.0
 
-        Returns
-        -------
-        score: float
-            Average Precision (AP_dsb) or Accuracy
+    Returns
+    -------
+    score: float
+        Average Precision (AP_dsb) or Accuracy
 
-        """
+    """
 
     gt_ids = np.arange(gt_image.shape[0])
     prediction_ids = np.unique(prediction_image)[1:]  # ignore background
@@ -288,8 +348,8 @@ def obtain_APdsb_one_hot(gt_image, prediction_image, ap_val):
 
     for j in range(iou_table.shape[0]):
         for k in range(iou_table.shape[1]):
-            intersection = ((gt_image[k] > 0) & (prediction_image == prediction_ids[j]))
-            union = ((gt_image[k] > 0) | (prediction_image == prediction_ids[j]))
+            intersection = (gt_image[k] > 0) & (prediction_image == prediction_ids[j])
+            union = (gt_image[k] > 0) | (prediction_image == prediction_ids[j])
             iou_table[j, k] = np.sum(intersection) / np.sum(union)
 
     iou_table_binary = iou_table >= ap_val
diff --git a/EmbedSeg/utils/preprocess_data.py b/EmbedSeg/utils/preprocess_data.py
index 9cd920e..6df794e 100644
--- a/EmbedSeg/utils/preprocess_data.py
+++ b/EmbedSeg/utils/preprocess_data.py
@@ -10,71 +10,91 @@
 from tqdm import tqdm
 
 
-def extract_data(zip_url, project_name, data_dir='../../../data/'):
+def extract_data(zip_url, project_name, data_dir="../../../data/"):
     """
-        Extracts data from `zip_url` to the location identified by `data_dir` and `project_name` parameters.
-
-        Parameters
-        ----------
-        zip_url: string
-            Indicates the external url
-        project_name: string
-            Indicates the path to the sub-directory at the location identified by the parameter `data_dir`
-        data_dir: string
-            Indicates the path to the directory where the data should be saved.
-        Returns
-        -------
+    Extracts data from `zip_url` to the location identified by
+    `data_dir` and `project_name` parameters.
+
+    Parameters
+    ----------
+    zip_url: string
+        Indicates the external url
+    project_name: string
+        Indicates the path to the sub-directory at the location
+        identified by the parameter `data_dir`
+    data_dir: string
+        Indicates the path to the directory where the data should be saved.
+    Returns
+    -------
 
     """
-    zip_path = os.path.join(data_dir, project_name + '.zip')
+    zip_path = os.path.join(data_dir, project_name + ".zip")
 
     if not os.path.exists(data_dir):
         os.makedirs(data_dir)
         print("Created new directory {}".format(data_dir))
 
-    if (os.path.exists(zip_path)):
+    if os.path.exists(zip_path):
         print("Zip file was downloaded and extracted before!")
     else:
-        if (os.path.exists(os.path.join(data_dir, project_name, 'download/'))):
+        if os.path.exists(os.path.join(data_dir, project_name, "download/")):
             pass
         else:
-            os.makedirs(os.path.join(data_dir, project_name, 'download/'))
+            os.makedirs(os.path.join(data_dir, project_name, "download/"))
             urllib.request.urlretrieve(zip_url, zip_path)
             print("Downloaded data as {}".format(zip_path))
-            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+            with zipfile.ZipFile(zip_path, "r") as zip_ref:
                 zip_ref.extractall(data_dir)
-            if os.path.exists(os.path.join(data_dir, os.path.basename(zip_url)[:-4], 'train')):
-                shutil.move(os.path.join(data_dir, os.path.basename(zip_url)[:-4], 'train'),
-                            os.path.join(data_dir, project_name, 'download/'))
-            if os.path.exists(os.path.join(data_dir, os.path.basename(zip_url)[:-4], 'val')):
-                shutil.move(os.path.join(data_dir, os.path.basename(zip_url)[:-4], 'val'),
-                            os.path.join(data_dir, project_name, 'download/'))
-            if os.path.exists(os.path.join(data_dir, os.path.basename(zip_url)[:-4], 'test')):
-                shutil.move(os.path.join(data_dir, os.path.basename(zip_url)[:-4], 'test'),
-                            os.path.join(data_dir, project_name, 'download/'))
-            print("Unzipped data to {}".format(os.path.join(data_dir, project_name, 'download/')))
+            if os.path.exists(
+                os.path.join(data_dir, os.path.basename(zip_url)[:-4], "train")
+            ):
+                shutil.move(
+                    os.path.join(data_dir, os.path.basename(zip_url)[:-4], "train"),
+                    os.path.join(data_dir, project_name, "download/"),
+                )
+            if os.path.exists(
+                os.path.join(data_dir, os.path.basename(zip_url)[:-4], "val")
+            ):
+                shutil.move(
+                    os.path.join(data_dir, os.path.basename(zip_url)[:-4], "val"),
+                    os.path.join(data_dir, project_name, "download/"),
+                )
+            if os.path.exists(
+                os.path.join(data_dir, os.path.basename(zip_url)[:-4], "test")
+            ):
+                shutil.move(
+                    os.path.join(data_dir, os.path.basename(zip_url)[:-4], "test"),
+                    os.path.join(data_dir, project_name, "download/"),
+                )
+            print(
+                "Unzipped data to {}".format(
+                    os.path.join(data_dir, project_name, "download/")
+                )
+            )
 
 
 def make_dirs(data_dir, project_name):
     """
-        Makes directories - `train`, `val, `test` and subdirectories under each `images` and `masks`
+    Makes directories - `train`, `val, `test` and subdirectories
+    under each `images` and `masks`
 
-        Parameters
-        ----------
-        data_dir: string
-            Indicates the path where the `project` lives.
-        project_name: string
-            Indicates the name of the sub-folder under the location identified by `data_dir`.
+    Parameters
+    ----------
+    data_dir: string
+        Indicates the path where the `project` lives.
+    project_name: string
+        Indicates the name of the sub-folder under the location
+        identified by `data_dir`.
 
-        Returns
-        -------
+    Returns
+    -------
     """
-    image_path_train = os.path.join(data_dir, project_name, 'train', 'images/')
-    instance_path_train = os.path.join(data_dir, project_name, 'train', 'masks/')
-    image_path_val = os.path.join(data_dir, project_name, 'val', 'images/')
-    instance_path_val = os.path.join(data_dir, project_name, 'val', 'masks/')
-    image_path_test = os.path.join(data_dir, project_name, 'test', 'images/')
-    instance_path_test = os.path.join(data_dir, project_name, 'test', 'masks/')
+    image_path_train = os.path.join(data_dir, project_name, "train", "images/")
+    instance_path_train = os.path.join(data_dir, project_name, "train", "masks/")
+    image_path_val = os.path.join(data_dir, project_name, "val", "images/")
+    instance_path_val = os.path.join(data_dir, project_name, "val", "masks/")
+    image_path_test = os.path.join(data_dir, project_name, "test", "images/")
+    instance_path_test = os.path.join(data_dir, project_name, "test", "masks/")
 
     if not os.path.exists(image_path_train):
         os.makedirs(os.path.dirname(image_path_train))
@@ -101,10 +121,18 @@ def make_dirs(data_dir, project_name):
         print("Created new directory : {}".format(instance_path_test))
 
 
-def split_train_crops(project_name, center, crops_dir='crops', subset=0.15, by_fraction=True, train_name='train',
-                      seed=1000):
+def split_train_crops(
+    project_name,
+    center,
+    crops_dir="crops",
+    subset=0.15,
+    by_fraction=True,
+    train_name="train",
+    seed=1000,
+):
     """
-    Moves a fraction of crops (specified by `subset`) extracted from train images into a separate directory
+    Moves a fraction of crops (specified by `subset`) extracted from
+    train images into a separate directory
 
     Parameters
     ----------
@@ -115,11 +143,14 @@ def split_train_crops(project_name, center, crops_dir='crops', subset=0.15, by_f
     crops_dir: string
             Name of the crops directory. Default value = 'crops'
     subset: int/float
-            if by_fraction is True, then subset should be set equal to the percentage of image crops reserved for validation.
-            if by_fraction is False, then subset should be set equal to the number of image crops reserved for validation.
+            if by_fraction is True, then subset should be set equal to the
+            percentage of image crops reserved for validation.
+            if by_fraction is False, then subset should be set equal to the
+            number of image crops reserved for validation.
     by_fraction: boolean
             if True, then reserve a fraction <1 of image crops for validation
-            if False, then reserve the absolute number of image crops as specified by `subset` for validation
+            if False, then reserve the absolute number of image crops
+            as specified by `subset` for validation
     train_name: string
             name of directory containing train image and instance crops
     seed: int
@@ -128,28 +159,34 @@ def split_train_crops(project_name, center, crops_dir='crops', subset=0.15, by_f
     -------
     """
 
-    image_dir = os.path.join(crops_dir, project_name, train_name, 'images')
-    instance_dir = os.path.join(crops_dir, project_name, train_name, 'masks')
-    center_dir = os.path.join(crops_dir, project_name, train_name, 'center-' + center)
+    image_dir = os.path.join(crops_dir, project_name, train_name, "images")
+    instance_dir = os.path.join(crops_dir, project_name, train_name, "masks")
+    center_dir = os.path.join(crops_dir, project_name, train_name, "center-" + center)
 
-    image_names = sorted(glob(os.path.join(image_dir, '*.tif')))
-    instance_names = sorted(glob(os.path.join(instance_dir, '*')))  # this could be `tifs` or `csvs`
-    center_names = sorted(glob(os.path.join(center_dir, '*')))  # this could be `tifs` or `csvs`
+    image_names = sorted(glob(os.path.join(image_dir, "*.tif")))
+    instance_names = sorted(
+        glob(os.path.join(instance_dir, "*"))
+    )  # this could be `tifs` or `csvs`
+    center_names = sorted(
+        glob(os.path.join(center_dir, "*"))
+    )  # this could be `tifs` or `csvs`
 
     indices = np.arange(len(image_names))
     np.random.seed(seed)
     np.random.shuffle(indices)
 
-    if (by_fraction):
+    if by_fraction:
         subset_len = int(subset * len(image_names))
     else:
         subset_len = int(subset)
 
     valIndices = indices[:subset_len]
 
-    image_path_val = os.path.join(crops_dir, project_name, 'val', 'images/')
-    instance_path_val = os.path.join(crops_dir, project_name, 'val', 'masks/')
-    center_path_val = os.path.join(crops_dir, project_name, 'val', 'center-' + center + '/')
+    image_path_val = os.path.join(crops_dir, project_name, "val", "images/")
+    instance_path_val = os.path.join(crops_dir, project_name, "val", "masks/")
+    center_path_val = os.path.join(
+        crops_dir, project_name, "val", "center-" + center + "/"
+    )
 
     val_images_exist = False
     val_masks_exist = False
@@ -175,48 +212,72 @@ def split_train_crops(project_name, center, crops_dir='crops', subset=0.15, by_f
 
     if not val_images_exist and not val_masks_exist and not val_center_images_exist:
         for val_index in valIndices:
-            shutil.move(image_names[val_index], os.path.join(crops_dir, project_name, 'val', 'images'))
-            shutil.move(instance_names[val_index], os.path.join(crops_dir, project_name, 'val', 'masks'))
-            shutil.move(center_names[val_index], os.path.join(crops_dir, project_name, 'val', 'center-' + center))
+            shutil.move(
+                image_names[val_index],
+                os.path.join(crops_dir, project_name, "val", "images"),
+            )
+            shutil.move(
+                instance_names[val_index],
+                os.path.join(crops_dir, project_name, "val", "masks"),
+            )
+            shutil.move(
+                center_names[val_index],
+                os.path.join(crops_dir, project_name, "val", "center-" + center),
+            )
 
-        print("Val Images/Masks/Center-{}-image crops saved at {}".format(center,
-                                                                          os.path.join(crops_dir, project_name, 'val')))
+        print(
+            "Val Images/Masks/Center-{}-image crops saved at {}".format(
+                center, os.path.join(crops_dir, project_name, "val")
+            )
+        )
     else:
-        print("Val Images/Masks/Center-{}-image crops already available at {}".format(center, os.path.join(crops_dir,
-                                                                                                           project_name,
-                                                                                                           'val')))
+        print(
+            "Val Images/Masks/Center-{}-image crops already available at {}".format(
+                center, os.path.join(crops_dir, project_name, "val")
+            )
+        )
 
 
-def split_train_val(data_dir, project_name, train_val_name, subset=0.15, by_fraction=True, seed=1000):
+def split_train_val(
+    data_dir, project_name, train_val_name, subset=0.15, by_fraction=True, seed=1000
+):
     """
-        Splits the `train` directory into `val` directory using the partition percentage of `subset`.
-
-        Parameters
-        ----------
-        data_dir: string
-            Indicates the path where the `project` lives.
-        project_name: string
-            Indicates the name of the sub-folder under the location identified by `data_dir`.
-        train_val_name: string
-            Indicates the name of the sub-directory under `project-name` which must be split
-        subset: float
-            Indicates the fraction of data to be reserved for validation
-        seed: integer
-            Allows for the same partition to be used in each experiment.
-            Change this if you would like to obtain results with different train-val partitions.
-        Returns
-        -------
+    Splits the `train` directory into `val` directory
+    using the partition percentage of `subset`.
+
+    Parameters
+    ----------
+    data_dir: string
+        Indicates the path where the `project` lives.
+    project_name: string
+        Indicates the name of the sub-folder under the location
+        identified by `data_dir`.
+    train_val_name: string
+        Indicates the name of the sub-directory under `project-name`
+        which must be split
+    subset: float
+        Indicates the fraction of data to be reserved for validation
+    seed: integer
+        Allows for the same partition to be used in each experiment.
+        Change this if you would like to obtain results with different
+        train-val partitions.
+    Returns
+    -------
 
     """
 
-    image_dir = os.path.join(data_dir, project_name, 'download', train_val_name, 'images')
-    instance_dir = os.path.join(data_dir, project_name, 'download', train_val_name, 'masks')
-    image_names = sorted(glob(os.path.join(image_dir, '*.tif')))
-    instance_names = sorted(glob(os.path.join(instance_dir, '*.tif')))
+    image_dir = os.path.join(
+        data_dir, project_name, "download", train_val_name, "images"
+    )
+    instance_dir = os.path.join(
+        data_dir, project_name, "download", train_val_name, "masks"
+    )
+    image_names = sorted(glob(os.path.join(image_dir, "*.tif")))
+    instance_names = sorted(glob(os.path.join(instance_dir, "*.tif")))
     indices = np.arange(len(image_names))
     np.random.seed(seed)
     np.random.shuffle(indices)
-    if (by_fraction):
+    if by_fraction:
         subset_len = int(subset * len(image_names))
     else:
         subset_len = int(subset)
@@ -225,54 +286,87 @@ def split_train_val(data_dir, project_name, train_val_name, subset=0.15, by_frac
     make_dirs(data_dir=data_dir, project_name=project_name)
 
     for val_index in val_indices:
-        shutil.copy(image_names[val_index], os.path.join(data_dir, project_name, 'val', 'images'))
-        shutil.copy(instance_names[val_index], os.path.join(data_dir, project_name, 'val', 'masks'))
+        shutil.copy(
+            image_names[val_index],
+            os.path.join(data_dir, project_name, "val", "images"),
+        )
+        shutil.copy(
+            instance_names[val_index],
+            os.path.join(data_dir, project_name, "val", "masks"),
+        )
 
     for trainIndex in trainIndices:
-        shutil.copy(image_names[trainIndex], os.path.join(data_dir, project_name, 'train', 'images'))
-        shutil.copy(instance_names[trainIndex], os.path.join(data_dir, project_name, 'train', 'masks'))
-
-    image_dir = os.path.join(data_dir, project_name, 'download', 'test', 'images')
-    instance_dir = os.path.join(data_dir, project_name, 'download', 'test', 'masks')
-    image_names = sorted(glob(os.path.join(image_dir, '*.tif')))
-    instance_names = sorted(glob(os.path.join(instance_dir, '*.tif')))
+        shutil.copy(
+            image_names[trainIndex],
+            os.path.join(data_dir, project_name, "train", "images"),
+        )
+        shutil.copy(
+            instance_names[trainIndex],
+            os.path.join(data_dir, project_name, "train", "masks"),
+        )
+
+    image_dir = os.path.join(data_dir, project_name, "download", "test", "images")
+    instance_dir = os.path.join(data_dir, project_name, "download", "test", "masks")
+    image_names = sorted(glob(os.path.join(image_dir, "*.tif")))
+    instance_names = sorted(glob(os.path.join(instance_dir, "*.tif")))
     test_indices = np.arange(len(image_names))
     for test_index in test_indices:
-        shutil.copy(image_names[test_index], os.path.join(data_dir, project_name, 'test', 'images'))
-        shutil.copy(instance_names[test_index], os.path.join(data_dir, project_name, 'test', 'masks'))
-    print("Train-Val-Test Images/Masks copied to {}".format(os.path.join(data_dir, project_name)))
+        shutil.copy(
+            image_names[test_index],
+            os.path.join(data_dir, project_name, "test", "images"),
+        )
+        shutil.copy(
+            instance_names[test_index],
+            os.path.join(data_dir, project_name, "test", "masks"),
+        )
+    print(
+        "Train-Val-Test Images/Masks copied to {}".format(
+            os.path.join(data_dir, project_name)
+        )
+    )
+
+
+def split_train_test(
+    data_dir, project_name, train_test_name, subset=0.5, by_fraction=True, seed=1000
+):
+    """
+    Splits the `train` directory into `test` directory
+    using the partition percentage of `subset`.
 
+    Parameters
+    ----------
+    data_dir: string
+        Indicates the path where the `project` lives.
+    project_name: string
+        Indicates the name of the sub-folder under the location
+        identified by `data_dir`.
+    train_test_name: string
+        Indicates the name of the sub-directory under `project-name`
+        which must be split
+    subset: float
+        Indicates the fraction of data to be reserved for evaluation
+    seed: integer
+        Allows for the same partition to be used
+        in each experiment.
+        Change this if you would like to obtain results
+        with different train-test partitions.
 
-def split_train_test(data_dir, project_name, train_test_name, subset=0.5, by_fraction=True, seed=1000):
-    """
-        Splits the `train` directory into `test` directory using the partition percentage of `subset`.
-
-        Parameters
-        ----------
-        data_dir: string
-            Indicates the path where the `project` lives.
-        project_name: string
-            Indicates the name of the sub-folder under the location identified by `data_dir`.
-        train_test_name: string
-            Indicates the name of the sub-directory under `project-name` which must be split
-        subset: float
-            Indicates the fraction of data to be reserved for evaluation
-        seed: integer
-            Allows for the same partition to be used in each experiment.
-            Change this if you would like to obtain results with different train-test partitions.
-
-        Returns
-        -------
+    Returns
+    -------
     """
 
-    image_dir = os.path.join(data_dir, project_name, 'download', train_test_name, 'images')
-    instance_dir = os.path.join(data_dir, project_name, 'download', train_test_name, 'masks')
-    image_names = sorted(glob(os.path.join(image_dir, '*.tif')))
-    instance_names = sorted(glob(os.path.join(instance_dir, '*.tif')))
+    image_dir = os.path.join(
+        data_dir, project_name, "download", train_test_name, "images"
+    )
+    instance_dir = os.path.join(
+        data_dir, project_name, "download", train_test_name, "masks"
+    )
+    image_names = sorted(glob(os.path.join(image_dir, "*.tif")))
+    instance_names = sorted(glob(os.path.join(instance_dir, "*.tif")))
     indices = np.arange(len(image_names))
     np.random.seed(seed)
     np.random.shuffle(indices)
-    if (by_fraction):
+    if by_fraction:
         subset_len = int(subset * len(image_names))
     else:
         subset_len = int(subset)
@@ -280,27 +374,54 @@ def split_train_test(data_dir, project_name, train_test_name, subset=0.5, by_fra
     # make_dirs(data_dir=data_dir, project_name=project_name)
     test_images_exist = False
     test_masks_exist = False
-    if not os.path.exists(os.path.join(data_dir, project_name, 'download', 'test', 'images')):
-        os.makedirs(os.path.join(data_dir, project_name, 'download', 'test', 'images'))
-        print("Created new directory : {}".format(os.path.join(data_dir, project_name, 'download', 'test', 'images')))
+    if not os.path.exists(
+        os.path.join(data_dir, project_name, "download", "test", "images")
+    ):
+        os.makedirs(os.path.join(data_dir, project_name, "download", "test", "images"))
+        print(
+            "Created new directory : {}".format(
+                os.path.join(data_dir, project_name, "download", "test", "images")
+            )
+        )
     else:
         test_images_exist = True
-    if not os.path.exists(os.path.join(data_dir, project_name, 'download', 'test', 'masks')):
-        os.makedirs(os.path.join(data_dir, project_name, 'download', 'test', 'masks'))
-        print("Created new directory : {}".format(os.path.join(data_dir, project_name, 'download', 'test', 'masks')))
+    if not os.path.exists(
+        os.path.join(data_dir, project_name, "download", "test", "masks")
+    ):
+        os.makedirs(os.path.join(data_dir, project_name, "download", "test", "masks"))
+        print(
+            "Created new directory : {}".format(
+                os.path.join(data_dir, project_name, "download", "test", "masks")
+            )
+        )
     else:
         test_masks_exist = True
     if not test_images_exist and not test_masks_exist:
         for test_index in test_indices:
-            shutil.move(image_names[test_index], os.path.join(data_dir, project_name, 'download', 'test', 'images'))
-            shutil.move(instance_names[test_index], os.path.join(data_dir, project_name, 'download', 'test', 'masks'))
-        print("Train-Test Images/Masks saved at {}".format(os.path.join(data_dir, project_name, 'download')))
+            shutil.move(
+                image_names[test_index],
+                os.path.join(data_dir, project_name, "download", "test", "images"),
+            )
+            shutil.move(
+                instance_names[test_index],
+                os.path.join(data_dir, project_name, "download", "test", "masks"),
+            )
+        print(
+            "Train-Test Images/Masks saved at {}".format(
+                os.path.join(data_dir, project_name, "download")
+            )
+        )
     else:
         print(
-            "Train-Test Images/Masks already available at {}".format(os.path.join(data_dir, project_name, 'download')))
+            "Train-Test Images/Masks already available at {}".format(
+                os.path.join(data_dir, project_name, "download")
+            )
+        )
 
 
-def calculate_foreground_weight(data_dir, project_name, train_val_name, mode, background_id=0):
+def calculate_foreground_weight(
+    data_dir, project_name, train_val_name, mode, background_id=0
+):
     """
 
     Parameters
@@ -320,30 +441,37 @@ def calculate_foreground_weight(data_dir, project_name, train_val_name, mode, ba
     Returns
     -------
     float:
-        Ratio of the number of foreground pixels to the background pixels, averaged over all available label masks
+        Ratio of the number of foreground pixels to the background pixels,
+        averaged over all available label masks
 
     """
     instance_names = []
     for name in train_val_name:
-        instance_dir = os.path.join(data_dir, project_name, name, 'masks')
-        instance_names += sorted(glob(os.path.join(instance_dir, '*.tif')))
+        instance_dir = os.path.join(data_dir, project_name, name, "masks")
+        instance_names += sorted(glob(os.path.join(instance_dir, "*.tif")))
 
     statistics = []
     for i in tqdm(range(len(instance_names)), position=0, leave=True):
         ma = tifffile.imread(instance_names[i])
-        if (mode in ['2d', '3d_sliced', '3d_ilp']):
+        if mode in ["2d", "3d_sliced", "3d_ilp"]:
             statistics.append(10.0)
-        elif (mode == '3d'):
+        elif mode == "3d":
             z, y, x = np.where(ma == background_id)
             len_bg = len(z)
             z, y, x = np.where(ma > background_id)
             len_fg = len(z)
             statistics.append(len_bg / len_fg)
-    print("Foreground weight of the `{}` dataset set equal to {:.3f}".format(project_name, np.mean(statistics)))
+    print(
+        "Foreground weight of the `{}` dataset set equal to {:.3f}".format(
+            project_name, np.mean(statistics)
+        )
+    )
     return np.mean(statistics)
 
 
-def calculate_object_size(data_dir, project_name, train_val_name, mode, one_hot, process_k, background_id=0):
+def calculate_object_size(
+    data_dir, project_name, train_val_name, mode, one_hot, process_k, background_id=0
+):
     """
     Calculate the mean object size from the available label masks
 
@@ -353,7 +481,8 @@ def calculate_object_size(data_dir, project_name, train_val_name, mode, one_hot,
     data_dir: string
         Name of directory storing the data. For example, 'data'
     project_name: string
-        Name of directory containing data specific to this project. For example, 'dsb-2018'
+        Name of directory containing data specific to this project.
+        For example, 'dsb-2018'
     train_val_name: string
         Name of directory containing 'train' and 'val' images and instance masks
     mode: string
@@ -361,19 +490,26 @@ def calculate_object_size(data_dir, project_name, train_val_name, mode, one_hot,
     one_hot: boolean
         set to True, if instances are encoded as one-hot
     process_k: tuple (int, int)
-        Parameter for speeding up the calculation of the object size by considering only a fewer number of images and objects
+        Parameter for speeding up the calculation of the object size
+        by considering only a fewer number of images and objects
         The first argument in the tuple is the number of images one should consider
-        The second argument in the tuple is the number of objects in the image, one should consider
+        The second argument in the tuple is the number of objects in the image,
+        one should consider
     background_id: int
         Id which corresponds to the background.
 
     Returns
     -------
     (float, float, float, float, float, float, float, float, float)
-    (minimum number of pixels in an object, mean number of pixels in an object, max number of pixels in an object,
-    mean number of pixels along the x dimension, standard deviation of number of pixels along the x dimension,
-    mean number of pixels along the y dimension, standard deviation of number of pixels along the y dimension,
-    mean number of pixels along the z dimension, standard deviation of number of pixels along the z dimension)
+    (minimum number of pixels in an object,
+     mean number of pixels in an object,
+     max number of pixels in an object,
+    mean number of pixels along the x dimension,
+     standard deviation of number of pixels along the x dimension,
+    mean number of pixels along the y dimension,
+     standard deviation of number of pixels along the y dimension,
+    mean number of pixels along the z dimension,
+     standard deviation of number of pixels along the z dimension)
 
     """
 
@@ -383,8 +519,8 @@ def calculate_object_size(data_dir, project_name, train_val_name, mode, one_hot,
     size_list_z = []
     size_list = []
     for name in train_val_name:
-        instance_dir = os.path.join(data_dir, project_name, name, 'masks')
-        instance_names += sorted(glob(os.path.join(instance_dir, '*.tif')))
+        instance_dir = os.path.join(data_dir, project_name, name, "masks")
+        instance_names += sorted(glob(os.path.join(instance_dir, "*.tif")))
 
     if process_k is not None:
         n_images = process_k[0]
@@ -392,13 +528,13 @@ def calculate_object_size(data_dir, project_name, train_val_name, mode, one_hot,
         n_images = len((instance_names))
     for i in tqdm(range(len(instance_names[:n_images])), position=0, leave=True):
         ma = tifffile.imread(instance_names[i])
-        if (one_hot and mode == '2d'):
+        if one_hot and mode == "2d":
             for z in range(ma.shape[0]):
                 y, x = np.where(ma[z] == 1)
                 size_list_x.append(np.max(x) - np.min(x))
                 size_list_y.append(np.max(y) - np.min(y))
                 size_list.append(len(x))
-        elif (not one_hot and mode == '2d'):
+        elif not one_hot and mode == "2d":
             ids = np.unique(ma)
             ids = ids[ids != background_id]
             for id in ids:
@@ -406,7 +542,7 @@ def calculate_object_size(data_dir, project_name, train_val_name, mode, one_hot,
                 size_list_x.append(np.max(x) - np.min(x))
                 size_list_y.append(np.max(y) - np.min(y))
                 size_list.append(len(x))
-        elif (not one_hot and mode in ['3d', '3d_sliced']):
+        elif not one_hot and mode in ["3d", "3d_sliced"]:
             ids = np.unique(ma)
             ids = ids[ids != background_id]
             if process_k is not None:
@@ -420,7 +556,7 @@ def calculate_object_size(data_dir, project_name, train_val_name, mode, one_hot,
                 size_list_y.append(np.max(y) - np.min(y))
                 size_list_x.append(np.max(x) - np.min(x))
                 size_list.append(len(x))
-        elif (not one_hot and mode == '3d_ilp'):
+        elif not one_hot and mode == "3d_ilp":
             for z in range(ma.shape[0]):
                 ids = np.unique(ma[z])
                 ids = ids[ids != background_id]
@@ -430,47 +566,91 @@ def calculate_object_size(data_dir, project_name, train_val_name, mode, one_hot,
                     size_list_x.append(np.max(x) - np.min(x))
                     size_list.append(len(x))
 
-    print("Minimum object size of the `{}` dataset is equal to {}".format(project_name, np.min(size_list)))
-    print("Mean object size of the `{}` dataset is equal to {}".format(project_name, np.mean(size_list)))
-    print("Maximum object size of the `{}` dataset is equal to {}".format(project_name, np.max(size_list)))
-    print("Average object size of the `{}` dataset along `x` is equal to {:.3f}".format(project_name,
-                                                                                        np.mean(size_list_x)))
-    print("Std. dev object size of the `{}` dataset along `x` is equal to {:.3f}".format(project_name,
-                                                                                         np.std(size_list_x)))
-    print("Average object size of the `{}` dataset along `y` is equal to {:.3f}".format(project_name,
-                                                                                        np.mean(size_list_y)))
-    print("Std. dev object size of the `{}` dataset along `y` is equal to {:.3f}".format(project_name,
-                                                                                         np.std(size_list_y)))
-
-    if mode in ['3d', '3d_sliced']:
-        print("Average object size of the `{}` dataset along `z` is equal to {:.3f}".format(project_name,
-                                                                                            np.mean(size_list_z)))
-        print("Std. dev object size of the `{}` dataset along `z` is equal to {:.3f}".format(project_name,
-                                                                                             np.std(size_list_z)))
-        return np.min(size_list).astype(np.float), np.mean(size_list).astype(np.float), np.max(size_list).astype(
-            np.float), \
-               np.mean(size_list_z).astype(np.float), np.mean(size_list_y).astype(
-            np.float), np.mean(size_list_x).astype(np.float), \
-               np.std(size_list_z).astype(np.float), np.std(size_list_y).astype(np.float), np.std(size_list_x).astype(
-            np.float)
+    print(
+        "Minimum object size of the `{}` dataset is equal to {}".format(
+            project_name, np.min(size_list)
+        )
+    )
+    print(
+        "Mean object size of the `{}` dataset is equal to {}".format(
+            project_name, np.mean(size_list)
+        )
+    )
+    print(
+        "Maximum object size of the `{}` dataset is equal to {}".format(
+            project_name, np.max(size_list)
+        )
+    )
+    print(
+        "Average object size of the `{}` dataset along `x` = {:.3f}".format(
+            project_name, np.mean(size_list_x)
+        )
+    )
+    print(
+        "Std. dev object size of the `{}` dataset along `x` = {:.3f}".format(
+            project_name, np.std(size_list_x)
+        )
+    )
+    print(
+        "Average object size of the `{}` dataset along `y` = {:.3f}".format(
+            project_name, np.mean(size_list_y)
+        )
+    )
+    print(
+        "Std. dev object size of the `{}` dataset along `y` = {:.3f}".format(
+            project_name, np.std(size_list_y)
+        )
+    )
+
+    if mode in ["3d", "3d_sliced"]:
+        print(
+            "Average object size of the `{}` dataset along `z` = {:.3f}".format(
+                project_name, np.mean(size_list_z)
+            )
+        )
+        print(
+            "Std. dev object size of the `{}` dataset along `z` =  {:.3f}".format(
+                project_name, np.std(size_list_z)
+            )
+        )
+        return (
+            np.min(size_list).astype(float),
+            np.mean(size_list).astype(float),
+            np.max(size_list).astype(float),
+            np.mean(size_list_z).astype(float),
+            np.mean(size_list_y).astype(float),
+            np.mean(size_list_x).astype(float),
+            np.std(size_list_z).astype(float),
+            np.std(size_list_y).astype(float),
+            np.std(size_list_x).astype(float),
+        )
 
     else:
-        return np.min(size_list).astype(np.float), np.mean(size_list).astype(np.float), np.max(size_list).astype(
-            np.float), \
-               None, np.mean(size_list_y).astype(np.float), np.mean(size_list_x).astype(np.float), None, np.std(
-            size_list_y).astype(
-            np.float), np.std(size_list_x).astype(np.float)
+        return (
+            np.min(size_list).astype(float),
+            np.mean(size_list).astype(float),
+            np.max(size_list).astype(float),
+            None,
+            np.mean(size_list_y).astype(float),
+            np.mean(size_list_x).astype(float),
+            None,
+            np.std(size_list_y).astype(float),
+            np.std(size_list_x).astype(float),
+        )
 
 
 def get_gpu_memory():
     """
-        Identifies the max memory on the operating GPU
-        https://stackoverflow.com/questions/59567226/how-to-programmatically-determine-available-gpu-memory-with-tensorflow/59571639#59571639
+    Identifies the max memory on the operating GPU
+    https://stackoverflow.com/questions/59567226/how-to-programmatically
+    -determine-available-gpu-memory-with-tensorflow/59571639#59571639
 
     """
 
     command = "nvidia-smi --query-gpu=memory.total --format=csv"
-    memory_total_info = sp.check_output(command.split()).decode('ascii').split('\n')[:-1][1:]
+    memory_total_info = (
+        sp.check_output(command.split()).decode("ascii").split("\n")[:-1][1:]
+    )
     memory_total_values = [int(x.split()[0]) for i, x in enumerate(memory_total_info)]
     return memory_total_values
 
@@ -494,76 +674,80 @@ def round_up_8(x):
     return (x.astype(int) + 7) & (-8)
 
 
-def calculate_max_eval_image_size(data_dir, project_name, test_name, mode, anisotropy_factor=1.0,
-                                  scale_factor=4.0):
+def calculate_max_eval_image_size(
+    data_dir, project_name, test_name, mode, anisotropy_factor=1.0, scale_factor=4.0
+):
     """
-        Identifies the tile size to be used during training and evaluation.
-        We look for the largest evaluation image.
-        If the entire image could fit on the available GPU memory (based on an empirical idea), then  the tile size is set equal to those dimensions.
-        If the dimensions are larger, then a smaller tile size as specified by the GPU memory is used.
-
-        Parameters
-        -------
-
-        data_dir: string
-            Name of directory storing the data. For example, 'data'
-        project_name: string
-            Name of directory containing data specific to this project. For example, 'dsb-2018'
-        train_val_name: string
-            Name of directory containing 'train' and 'val' images and instance masks
-        mode: string
-            One of '2d', '3d', '3d_sliced', '3d_ilp'
-        anisotropy_factor: float
-            Ratio of the pixel size along the z dimension to the pixel size along the x/y dimension
-            Here we assume that the pixel size along the x and y dimension is the same
-        scale_factor: float, default
-            Used to evaluate the maximum GPU memory which shall be used during evaluation
+    Identifies the tile size to be used during training and evaluation.
+    We look for the largest evaluation image.
+    If the entire image could fit on the available GPU memory
+    (based on an empirical idea), then  the tile size is
+    set equal to those dimensions.
+    If the dimensions are larger, then a smaller tile size
+    as specified by the GPU memory is used.
 
-        Returns
-        -------
+    Parameters
+    -------
 
-        (int, int, int)
-        (tile size along z, tile size along y, tile size along x)
+    data_dir: string
+        Name of directory storing the data. For example, 'data'
+    project_name: string
+        Name of directory containing data specific to this project.
+        For example, 'dsb-2018'
+    train_val_name: string
+        Name of directory containing 'train' and 'val' images and instance masks
+    mode: string
+        One of '2d', '3d', '3d_sliced', '3d_ilp'
+    anisotropy_factor: float
+        Ratio of the pixel size along the z dimension to the pixel size
+        along the x/y dimension
+        Here we assume that the pixel size along the x and y dimension is the same
+    scale_factor: float, default
+        Used to evaluate the maximum GPU memory which
+        shall be used during evaluation
 
-        """
+    Returns
+    -------
+
+    (int, int, int)
+    (tile size along z, tile size along y, tile size along x)
 
+    """
 
     image_names = []
     size_z_list = []
     size_y_list = []
     size_x_list = []
     for name in test_name:
-        instance_dir = os.path.join(data_dir, project_name, name, 'images')
-        image_names += sorted(glob(os.path.join(instance_dir, '*.tif')))
+        instance_dir = os.path.join(data_dir, project_name, name, "images")
+        image_names += sorted(glob(os.path.join(instance_dir, "*.tif")))
 
     for i in tqdm(range(len(image_names)), position=0, leave=True):
         im = tifffile.imread(image_names[i])
-        if (mode == '2d'):
+        if mode == "2d":
             size_y_list.append(im.shape[0])
             size_x_list.append(im.shape[1])
-        elif (mode in ['3d', '3d_sliced', '3d_ilp']):
+        elif mode in ["3d", "3d_sliced", "3d_ilp"]:
             size_z_list.append(im.shape[0])
             size_y_list.append(im.shape[1])
             size_x_list.append(im.shape[2])
-    if mode in ['2d', '3d_ilp', '3d_sliced']:
+    if mode in ["2d", "3d_ilp", "3d_sliced"]:
         max_y = np.max(size_y_list)
         max_x = np.max(size_x_list)
         max_y, max_x = round_up_8(max_y), round_up_8(max_x)
         max_x_y = np.maximum(max_x, max_y)
 
-        total_mem = get_gpu_memory()[0] * 1e6  # Note: get_gpu_memory returns a list
-        tile_size_temp = np.asarray((total_mem / (2 * 4 * scale_factor)) ** (1 / 2))  # 2D
-
-        if tile_size_temp < max_x_y:
-            max_x, max_y = round_up_8(tile_size_temp), round_up_8(tile_size_temp)
-        else:
-            max_x, max_y = max_x_y, max_x_y
-        print("Tile size of the `{}` dataset set equal to ({}, {})".format(project_name, max_y, max_x))
-        if mode == '3d_sliced':
-            return max_y.astype(np.float), max_y.astype(np.float), max_x.astype(np.float)
+        max_x, max_y = max_x_y, max_x_y
+        print(
+            "Tile size of the `{}` dataset set equal to ({}, {})".format(
+                project_name, max_y, max_x
+            )
+        )
+        if mode == "3d_sliced":
+            return max_y.astype(float), max_y.astype(float), max_x.astype(float)
         else:
-            return None, max_y.astype(np.float), max_x.astype(np.float)
-    elif mode == '3d':
+            return None, max_y.astype(float), max_x.astype(float)
+    elif mode == "3d":
         max_z = np.max(size_z_list)
         max_y = np.max(size_y_list)
         max_x = np.max(size_x_list)
@@ -574,25 +758,25 @@ def calculate_max_eval_image_size(data_dir, project_name, test_name, mode, aniso
 
         max_x_y = np.maximum(max_x, max_y)
 
-        total_mem = get_gpu_memory()[0] * 1e6  # Note: get_gpu_memory returns a list
-        tile_size_temp = np.asarray((total_mem * anisotropy_factor / (3 * 4 * scale_factor)) ** (1 / 3))  # 3D
-        if (tile_size_temp ** 3) / anisotropy_factor < (max_x_y ** 2) * max_z:
-            max_x = round_up_8(tile_size_temp)
-            max_y = round_up_8(tile_size_temp)
-            max_z = round_up_8(tile_size_temp / anisotropy_factor)
-        else:
-            max_x = max_x_y
-            max_y = max_x_y
-            max_z = max_z
+        max_x = max_x_y
+        max_y = max_x_y
+        max_z = max_z
 
-        print("Tile size of the `{}` dataset set equal to  (n_z = {}, n_y = {}, n_x = {})".format(
-            project_name, max_z, max_y, max_x))
-        return max_z.astype(np.float), max_y.astype(np.float), max_x.astype(np.float)
+        print(
+            "Tile size of the `{}` dataset set equal to  \
+                    (n_z = {}, n_y = {}, n_x = {})".format(
+                project_name, max_z, max_y, max_x
+            )
+        )
+        return max_z.astype(float), max_y.astype(float), max_x.astype(float)
 
 
-def calculate_avg_background_intensity(data_dir, project_name, train_val_name, one_hot, background_id=0):
+def calculate_avg_background_intensity(
+    data_dir, project_name, train_val_name, one_hot, background_id=0
+):
     """
-    Calculates the average intensity in the regions of the raw image which corresponds to the background label
+    Calculates the average intensity in the regions of the raw image
+    which corresponds to the background label
 
     Parameters
     -------
@@ -617,10 +801,10 @@ def calculate_avg_background_intensity(data_dir, project_name, train_val_name, o
     instance_names = []
     image_names = []
     for name in train_val_name:
-        instance_dir = os.path.join(data_dir, project_name, name, 'masks')
-        image_dir = os.path.join(data_dir, project_name, name, 'images')
-        instance_names += sorted(glob(os.path.join(instance_dir, '*.tif')))
-        image_names += sorted(glob(os.path.join(image_dir, '*.tif')))
+        instance_dir = os.path.join(data_dir, project_name, name, "masks")
+        image_dir = os.path.join(data_dir, project_name, name, "images")
+        instance_names += sorted(glob(os.path.join(instance_dir, "*.tif")))
+        image_names += sorted(glob(os.path.join(image_dir, "*.tif")))
     statistics = []
     if one_hot:
         for i in tqdm(range(len(instance_names)), position=0, leave=True):
@@ -638,16 +822,33 @@ def calculate_avg_background_intensity(data_dir, project_name, train_val_name, o
             elif im.ndim == ma.ndim + 1:  # multi-channel image
                 statistics.append(np.average(im[:, bg_mask], axis=1))
     if im.ndim == ma.ndim:
-        print("Average background intensity of the `{}` dataset set equal to {:.3f}".format(project_name,
-                                                                                            np.mean(statistics, 0)))
+        print(
+            "Average background intensity of the `{}` dataset set \
+                    equal to {:.3f}".format(
+                project_name, np.mean(statistics, 0)
+            )
+        )
     elif im.ndim == ma.ndim + 1:
-        print("Average background intensity of the `{}` dataset set equal to {}".format(project_name,
-                                                                                        np.mean(statistics, 0)))
+        print(
+            "Average background intensity of the `{}` dataset set \
+                    equal to {}".format(
+                project_name, np.mean(statistics, 0)
+            )
+        )
     return np.mean(statistics, 0).tolist()
 
 
-def get_data_properties(data_dir, project_name, train_val_name, test_name, mode, one_hot=False, process_k=None,
-                        anisotropy_factor=1.0, background_id=0):
+def get_data_properties(
+    data_dir,
+    project_name,
+    train_val_name,
+    test_name,
+    mode,
+    one_hot=False,
+    process_k=None,
+    anisotropy_factor=1.0,
+    background_id=0,
+):
     """
 
     Parameters
@@ -666,36 +867,64 @@ def get_data_properties(data_dir, project_name, train_val_name, test_name, mode,
     one_hot: boolean
             set to True, if instances are encoded in a one-hot fashion
     process_k (int, int)
-            first `int` argument in tuple specifies number of images which must be processed
-            second `int` argument in tuple specifies number of ids which must be processed
+            first `int` argument in tuple specifies number of images
+            which must be processed
+            second `int` argument in tuple specifies number of ids
+            which must be processed
     anisotropy_factor: float
-            Ratio of the real-world size of the z-dimension to the x or y dimension in the raw images
-            If the image is down-sampled along the z-dimension, then `anisotropy_factor` is greater than 1.0
+            Ratio of the real-world size of the z-dimension to the
+            x or y dimension in the raw images
+            If the image is down-sampled along the z-dimension,
+            then `anisotropy_factor` is greater than 1.0
     background_id: int
             Label id corresponding to the background
 
     Returns
     -------
     data_properties_dir: dictionary
-            keys include `foreground_weight`, `min_object_size`, `project_name`, `avg_background_intensity` etc
+            keys include `foreground_weight`, `min_object_size`,
+            `project_name`, `avg_background_intensity` etc
 
     """
     data_properties_dir = {}
-    data_properties_dir['foreground_weight'] = calculate_foreground_weight(data_dir, project_name, train_val_name, mode,
-                                                                           background_id=background_id)
-    data_properties_dir['min_object_size'], data_properties_dir['mean_object_size'], data_properties_dir[
-        'max_object_size'], \
-    data_properties_dir['avg_object_size_z'], data_properties_dir[
-        'avg_object_size_y'], data_properties_dir['avg_object_size_x'], \
-    data_properties_dir['stdev_object_size_z'], data_properties_dir['stdev_object_size_y'], data_properties_dir[
-        'stdev_object_size_x'] = calculate_object_size(data_dir, project_name, train_val_name, mode, one_hot, process_k,
-                                                       background_id=background_id)
-    data_properties_dir['n_z'], data_properties_dir['n_y'], data_properties_dir['n_x'] = calculate_max_eval_image_size(
-        data_dir=data_dir, project_name=project_name, test_name=test_name, mode=mode,
-        anisotropy_factor=anisotropy_factor)
-    data_properties_dir['one_hot'] = one_hot
-    data_properties_dir['avg_background_intensity'] = calculate_avg_background_intensity(data_dir, project_name,
-                                                                                         train_val_name, one_hot,
-                                                                                         background_id=background_id)
-    data_properties_dir['project_name'] = project_name
+    data_properties_dir["foreground_weight"] = calculate_foreground_weight(
+        data_dir, project_name, train_val_name, mode, background_id=background_id
+    )
+    (
+        data_properties_dir["min_object_size"],
+        data_properties_dir["mean_object_size"],
+        data_properties_dir["max_object_size"],
+        data_properties_dir["avg_object_size_z"],
+        data_properties_dir["avg_object_size_y"],
+        data_properties_dir["avg_object_size_x"],
+        data_properties_dir["stdev_object_size_z"],
+        data_properties_dir["stdev_object_size_y"],
+        data_properties_dir["stdev_object_size_x"],
+    ) = calculate_object_size(
+        data_dir,
+        project_name,
+        train_val_name,
+        mode,
+        one_hot,
+        process_k,
+        background_id=background_id,
+    )
+    (
+        data_properties_dir["n_z"],
+        data_properties_dir["n_y"],
+        data_properties_dir["n_x"],
+    ) = calculate_max_eval_image_size(
+        data_dir=data_dir,
+        project_name=project_name,
+        test_name=test_name,
+        mode=mode,
+        anisotropy_factor=anisotropy_factor,
+    )
+    data_properties_dir["one_hot"] = one_hot
+    data_properties_dir[
+        "avg_background_intensity"
+    ] = calculate_avg_background_intensity(
+        data_dir, project_name, train_val_name, one_hot, background_id=background_id
+    )
+    data_properties_dir["project_name"] = project_name
     return data_properties_dir
diff --git a/EmbedSeg/utils/transforms.py b/EmbedSeg/utils/transforms.py
index 7ea7bad..6a8c1b2 100644
--- a/EmbedSeg/utils/transforms.py
+++ b/EmbedSeg/utils/transforms.py
@@ -7,7 +7,8 @@
 
 class RandomRotationsAndFlips(T.RandomRotation):
     """
-    A class used to represent Random Rotations and Flips for Augmenting 2D Image Data
+    A class used to represent Random Rotations and Flips for Augmenting
+    2D Image Data
 
     ...
 
@@ -17,11 +18,13 @@ class RandomRotationsAndFlips(T.RandomRotation):
         keys include `instance`, `label`, `center-image`
         See `TwoDimensionalDataset.py`
     one_hot : bool
-        Should be set to True, if the GT label masks are present in a one-hot encoded fashion
+        Should be set to True, if the GT label masks are present
+        in a one-hot encoded fashion
 
     Methods
     -------
-    __call__: Returns rotated or flipped image, instance label mask and center image
+    __call__: Returns rotated or flipped image, instance label mask and
+    center image
 
     """
 
@@ -33,7 +36,8 @@ def __init__(self, keys=[], one_hot=False, *args, **kwargs):
             keys include `instance`, `label`, `center-image`
             See `TwoDimensionalDataset.py`
         one_hot : bool
-            Should be set to True, if the GT label masks are present in a one-hot encoded fashion
+            Should be set to True, if the GT label masks are present in a
+            one-hot encoded fashion
 
         """
 
@@ -43,56 +47,63 @@ def __init__(self, keys=[], one_hot=False, *args, **kwargs):
 
     def __call__(self, sample):
         """
-            Parameters
-            ----------
-            sample
+        Parameters
+        ----------
+        sample
 
-            Returns
-            ----------
-            sample
+        Returns
+        ----------
+        sample
 
         """
 
-        angle = self.get_params(self.degrees)
+        self.get_params(self.degrees)
         times = np.random.choice(4)
         flip = np.random.choice(2)
 
         for idx, k in enumerate(self.keys):
-            assert (k in sample)
-            if (self.one_hot and k == 'instance'):
+            assert k in sample
+            if self.one_hot and k == "instance":
                 temp = np.ascontiguousarray(np.rot90(sample[k], times, (2, 3)))
             else:
                 temp = np.ascontiguousarray(np.rot90(sample[k], times, (1, 2)))
             if flip == 0:
                 sample[k] = temp
             else:
-                if (self.one_hot and k == 'instance'):
-                    sample[k] = np.ascontiguousarray(np.flip(temp, axis=2))  # flip about D - axis
+                if self.one_hot and k == "instance":
+                    sample[k] = np.ascontiguousarray(
+                        np.flip(temp, axis=2)
+                    )  # flip about D - axis
                 else:
-                    sample[k] = np.ascontiguousarray(np.flip(temp, axis=1))  # flip about Y - axis
+                    sample[k] = np.ascontiguousarray(
+                        np.flip(temp, axis=1)
+                    )  # flip about Y - axis
         return sample
 
 
 class RandomRotationsAndFlips_3d(T.RandomRotation):
     """
-        A class used to represent Random Rotations and Flips for Augmenting 3D Image Data
+    A class used to represent Random Rotations and Flips for Augmenting
+    3D Image Data
 
-        ...
+    ...
 
-        Attributes
-        ----------
-        keys : dictionary
-            keys include `instance`, `label`, `center-image`
-            See `ThreeDimensionalDataset.py`
-        one_hot : bool
-            Should be set to True, if the GT label masks are present in a one-hot encoded fashion
-            Not applicable to 3D. This parameter will be deprecated in a future release
+    Attributes
+    ----------
+    keys : dictionary
+        keys include `instance`, `label`, `center-image`
+        See `ThreeDimensionalDataset.py`
+    one_hot : bool
+        Should be set to True, if the GT label masks are present in a one-hot
+        encoded fashion
+        Not applicable to 3D. This parameter will be deprecated in a future release
 
-        Methods
-        -------
-        __call__: Returns rotated or flipped image, instance label mask and center image
+    Methods
+    -------
+    __call__: Returns rotated or flipped image, instance label mask and
+    center image
 
-        """
+    """
 
     def __init__(self, keys=[], one_hot=False, *args, **kwargs):
         """
@@ -102,8 +113,10 @@ def __init__(self, keys=[], one_hot=False, *args, **kwargs):
             keys include `instance`, `label`, `center-image`
             See `ThreeDimensionalDataset.py`
         one_hot : bool
-            Should be set to True, if the GT label masks are present in a one-hot encoded fashion
-            Not applicable to 3D. This parameter will be deprecated in a future release
+            Should be set to True, if the GT label masks are present
+            in a one-hot encoded fashion
+            Not applicable to 3D. This parameter will be deprecated
+            in a future release
 
         """
 
@@ -122,15 +135,14 @@ def __call__(self, sample):
         sample
 
         """
-        angle = self.get_params(self.degrees)
+        self.get_params(self.degrees)
         times = np.random.choice(4)
         flip = np.random.choice(2)
         dir_rot = np.random.choice(3)
         dir_flip = np.random.choice(3)
 
         for idx, k in enumerate(self.keys):
-
-            assert (k in sample)
+            assert k in sample
             if dir_rot == 0:  # rotate about ZY
                 temp = np.ascontiguousarray(np.rot90(sample[k], 2 * times, (1, 2)))
             elif dir_rot == 1:  # rotate about YX
@@ -153,47 +165,51 @@ def __call__(self, sample):
 
 class ToTensorFromNumpy(object):
     """
-        A class used to convert numpy arrays to PyTorch tensors
+    A class used to convert numpy arrays to PyTorch tensors
 
-        ...
+    ...
 
-        Attributes
-        ----------
-        keys : dictionary
-            keys include `instance`, `label`, `center-image`, `image`
-        type : str
+    Attributes
+    ----------
+    keys : dictionary
+        keys include `instance`, `label`, `center-image`, `image`
+    type : str
 
-        normalization_factor: float
+    normalization_factor: float
 
-        Methods
-        -------
-        __call__: Returns Pytorch Tensors
+    Methods
+    -------
+    __call__: Returns Pytorch Tensors
 
-            """
+    """
 
     def __init__(self, keys=[], type="float", normalization_factor=1.0):
-
-        if isinstance(type, collections.Iterable):
-            assert (len(keys) == len(type))
+        if isinstance(type, collections.abc.Iterable):
+            assert len(keys) == len(type)
 
         self.keys = keys
         self.type = type
         self.normalization_factor = normalization_factor
 
     def __call__(self, sample):
-
         for idx, k in enumerate(self.keys):
             # assert (k in sample)
 
             t = self.type
-            if isinstance(t, collections.Iterable):
+            if isinstance(t, collections.abc.Iterable):
                 t = t[idx]
-            if (k in sample):
-                if k == 'image':  # image
-                    sample[k] = torch.from_numpy(sample[k].astype("float32")).float().div(self.normalization_factor)
-                elif k == 'instance' or k == 'label':
-                    sample[k] = torch.from_numpy(sample[k])  # np.int16 to torch.int16 or short (since np.uint16 can not be directly cast to torch uint16)
-                elif k == 'center-image':
+            if k in sample:
+                if k == "image":  # image
+                    sample[k] = (
+                        torch.from_numpy(sample[k].astype("float32"))
+                        .float()
+                        .div(self.normalization_factor)
+                    )
+                elif k == "instance" or k == "label":
+                    sample[k] = torch.from_numpy(
+                        sample[k]
+                    )  # np.int16 to torch.int16 or short
+                elif k == "center-image":
                     sample[k] = torch.from_numpy(sample[k])  # np.bool to torch.Bool
         return sample
 
@@ -202,8 +218,8 @@ def get_transform(transforms):
     transform_list = []
 
     for tr in transforms:
-        name = tr['name']
-        opts = tr['opts']
+        name = tr["name"]
+        opts = tr["opts"]
 
         transform_list.append(globals()[name](**opts))
 
diff --git a/EmbedSeg/utils/utils.py b/EmbedSeg/utils/utils.py
index e5f0c48..7099ce7 100644
--- a/EmbedSeg/utils/utils.py
+++ b/EmbedSeg/utils/utils.py
@@ -13,7 +13,6 @@
 
 
 class AverageMeter(object):
-
     def __init__(self, num_classes=1):
         self.num_classes = num_classes
         self.reset()
@@ -30,53 +29,87 @@ def update(self, val, cl=0):
             self.sum[cl] += val
             self.count[cl] += 1
             self.avg_per_class = [
-                x / y if x > 0 else 0 for x, y in zip(self.sum, self.count)]
+                x / y if x > 0 else 0 for x, y in zip(self.sum, self.count)
+            ]
             self.avg = sum(self.avg_per_class) / len(self.avg_per_class)
 
 
 class Visualizer:
-
     def __init__(self, keys, cmap):
-
         self.fig, (self.ax1, self.ax2) = plt.subplots(2, 2, figsize=(10, 10))
 
-        self.ax1[1].axis('off')
-        self.ax2[0].axis('off')
-        self.ax2[1].axis('off')
+        self.ax1[1].axis("off")
+        self.ax2[0].axis("off")
+        self.ax2[1].axis("off")
         self.cmap = cmap
 
-    def display(self, image, key, title, center_x=None, center_y=None, samples_x=None, samples_y=None,
-                sample_spatial_embedding_x=None, sample_spatial_embedding_y=None,
-                sigma_x=None, sigma_y=None, color_sample=None, color_embedding=None):
-
-        n_images = len(image) if isinstance(image, (list, tuple)) else 1
+    def display(
+        self,
+        image,
+        key,
+        title,
+        center_x=None,
+        center_y=None,
+        samples_x=None,
+        samples_y=None,
+        sample_spatial_embedding_x=None,
+        sample_spatial_embedding_y=None,
+        sigma_x=None,
+        sigma_y=None,
+        color_sample=None,
+        color_embedding=None,
+    ):
+        len(image) if isinstance(image, (list, tuple)) else 1
 
         self.fig.frameon = False
-        if key == 'image':
+        if key == "image":
             self.ax1[0].cla()
-            self.ax1[0].imshow(self.prepare_img(image), cmap='magma')
-            self.ax1[0].axis('off')
-        elif key == 'groundtruth':
-
+            self.ax1[0].imshow(self.prepare_img(image), cmap="magma")
+            self.ax1[0].axis("off")
+        elif key == "groundtruth":
             self.ax1[1].cla()
-            self.ax1[1].imshow(self.prepare_img(image), cmap=self.cmap, interpolation='None')
-            self.ax1[1].axis('off')
-        elif key == 'prediction':
+            self.ax1[1].imshow(
+                self.prepare_img(image), cmap=self.cmap, interpolation="None"
+            )
+            self.ax1[1].axis("off")
+        elif key == "prediction":
             self.ax2[1].cla()
-            self.ax2[1].imshow(self.prepare_img(image), cmap=self.cmap, interpolation='None')
-            self.ax2[1].axis('off')
-        elif key == 'center':
+            self.ax2[1].imshow(
+                self.prepare_img(image), cmap=self.cmap, interpolation="None"
+            )
+            self.ax2[1].axis("off")
+        elif key == "center":
             self.ax2[0].cla()
-            self.ax2[0].imshow(self.prepare_img(image), cmap='gray')
+            self.ax2[0].imshow(self.prepare_img(image), cmap="gray")
             for i in range(len(color_sample.items())):
-                self.ax2[0].plot(center_x[i + 1], center_y[i + 1], color=color_embedding[i + 1], marker='x')
-                self.ax2[0].scatter(samples_x[i + 1], samples_y[i + 1], color=color_sample[i + 1], marker='+')
-                self.ax2[0].scatter(sample_spatial_embedding_x[i + 1], sample_spatial_embedding_y[i + 1],
-                                    color=color_embedding[i + 1], marker='.')
-                ellipse = Ellipse((center_x[i + 1], center_y[i + 1]), width=sigma_x[i + 1], height=sigma_y[i + 1],
-                                  angle=0, color=color_embedding[i + 1], alpha=0.5)
+                self.ax2[0].plot(
+                    center_x[i + 1],
+                    center_y[i + 1],
+                    color=color_embedding[i + 1],
+                    marker="x",
+                )
+                self.ax2[0].scatter(
+                    samples_x[i + 1],
+                    samples_y[i + 1],
+                    color=color_sample[i + 1],
+                    marker="+",
+                )
+                self.ax2[0].scatter(
+                    sample_spatial_embedding_x[i + 1],
+                    sample_spatial_embedding_y[i + 1],
+                    color=color_embedding[i + 1],
+                    marker=".",
+                )
+                ellipse = Ellipse(
+                    (center_x[i + 1], center_y[i + 1]),
+                    width=sigma_x[i + 1],
+                    height=sigma_y[i + 1],
+                    angle=0,
+                    color=color_embedding[i + 1],
+                    alpha=0.5,
+                )
                 self.ax2[0].add_artist(ellipse)
-            self.ax2[0].axis('off')
+            self.ax2[0].axis("off")
         plt.tight_layout()
         plt.draw()
         self.mypause(0.0001)
@@ -97,7 +130,7 @@ def prepare_img(image):
 
     @staticmethod
     def mypause(interval):
-        backend = plt.rcParams['backend']
+        backend = plt.rcParams["backend"]
         if backend in matplotlib.rcsetup.interactive_bk:
             figManager = matplotlib._pylab_helpers.Gcf.get_active()
             if figManager is not None:
@@ -110,82 +143,110 @@ def mypause(interval):
 
 class Cluster_3d:
     """
-            A class used to cluster pixel embeddings in 3D
-
-            Attributes
-            ----------
-            xyzm :  float (3, W, D, H)
-                    pixel coordinates of tile /grid
-
-            one_hot : bool
-                    Should be set to True, if the GT label masks are present in a one-hot encoded fashion
-                    Not applicable to 3D. This parameter will be deprecated in a future release
-
-            grid_x: int
-                    Length (width) of grid / tile
-
-            grid_y: int
-                    Height of grid / tile
-
-            grid_z: int
-                    Depth of grid / tile
-
-            pixel_x: float
-                    if grid_x = 1000 and pixel_x = 1.0, then the pixel spacing along the x direction is pixel_x/(grid_x-1) = 1/999
-            pixel_y: float
-                    if grid_y = 1000 and pixel_y = 1.0, then the pixel spacing along the y direction is pixel_y/(grid_y-1) = 1/999
-            pixel_z: float
-                    if grid_z = 1000 and pixel_z = 1, then the pixel spacing along the z direction is pixel_z/(grid_z-1) = 1/999
-
-
-            Methods
-            -------
-            __init__: Initializes an object of class `Cluster_3d`
-
-            cluster_with_gt: use the predicted spatial embeddings from all pixels belonging to the GT label mask
-                        to identify the predicted cluster (used during training and validation)
-
-            cluster:    use the  predicted spatial embeddings from all pixels in the test image.
-                        Employs `fg_thresh` and `seed_thresh`
-            cluster_local_maxima: use the  predicted spatial embeddings from all pixels in the test image.
-                        Employs only `fg_thresh`
-            """
+    A class used to cluster pixel embeddings in 3D
+
+    Attributes
+    ----------
+    xyzm :  float (3, W, D, H)
+            pixel coordinates of tile /grid
+
+    one_hot : bool
+            Should be set to True, if the GT label masks are present
+            in a one-hot encoded fashion
+            Not applicable to 3D. This parameter will be deprecated
+            in a future release
+
+    grid_x: int
+            Length (width) of grid / tile
+
+    grid_y: int
+            Height of grid / tile
+
+    grid_z: int
+            Depth of grid / tile
+
+    pixel_x: float
+            if grid_x = 1000 and pixel_x = 1.0, then the pixel spacing
+            along the x direction is pixel_x/(grid_x-1) = 1/999
+    pixel_y: float
+            if grid_y = 1000 and pixel_y = 1.0, then the pixel spacing
+            along the y direction is pixel_y/(grid_y-1) = 1/999
+    pixel_z: float
+            if grid_z = 1000 and pixel_z = 1, then the pixel spacing
+            along the z direction is pixel_z/(grid_z-1) = 1/999
+
+
+    Methods
+    -------
+    __init__: Initializes an object of class `Cluster_3d`
+
+    cluster_with_gt: use the predicted spatial embeddings
+            from all pixels belonging to the GT label mask
+            to identify the predicted cluster
+            (used during training and validation)
+
+    cluster: use the  predicted spatial embeddings
+             from all pixels in the test image.
+             Employs `fg_thresh` and `seed_thresh`
+    cluster_local_maxima: use the  predicted spatial embeddings
+            from all pixels in the test image.
+            Employs only `fg_thresh`
+    """
 
-    def __init__(self, grid_z, grid_y, grid_x, pixel_z, pixel_y, pixel_x, one_hot=False):
+    def __init__(
+        self, grid_z, grid_y, grid_x, pixel_z, pixel_y, pixel_x, device, one_hot=False
+    ):
         """
-           Parameters
-           ----------
-           xyzm :  float (3, W, D, H)
-                    pixel coordinates of tile /grid
-
-            one_hot : bool
-                    Should be set to True, if the GT label masks are present in a one-hot encoded fashion
-                    Not applicable to 3D. This parameter will be deprecated in a future release
-
-            grid_x: int
-                    Length (width) of grid / tile
-
-            grid_y: int
-                    Height of grid / tile
-
-            grid_z: int
-                    Depth of grid / tile
-
-            pixel_x: float
-                    if grid_x = 1000 and pixel_x = 1.0, then the pixel spacing along the x direction is pixel_x/(grid_x-1) = 1/999
-            pixel_y: float
-                    if grid_y = 1000 and pixel_y = 1.0, then the pixel spacing along the y direction is pixel_y/(grid_y-1) = 1/999
-            pixel_z: float
-                    if grid_z = 1000 and pixel_z = 1, then the pixel spacing along the z direction is pixel_z/(grid_z-1) = 1/999
+        Parameters
+        ----------
+        xyzm :  float (3, W, D, H)
+                 pixel coordinates of tile /grid
+
+         one_hot : bool
+                 Should be set to True, if the GT label masks are present
+                 in a one-hot encoded fashion
+                 Not applicable to 3D. This parameter will be deprecated
+                 in a future release
+
+         grid_x: int
+                 Length (width) of grid / tile
+
+         grid_y: int
+                 Height of grid / tile
+
+         grid_z: int
+                 Depth of grid / tile
+
+         pixel_x: float
+                 if grid_x = 1000 and pixel_x = 1.0, then the pixel spacing
+                 along the x direction is pixel_x/(grid_x-1) = 1/999
+         pixel_y: float
+                 if grid_y = 1000 and pixel_y = 1.0, then the pixel spacing
+                 along the y direction is pixel_y/(grid_y-1) = 1/999
+         pixel_z: float
+                 if grid_z = 1000 and pixel_z = 1, then the pixel spacing
+                 along the z direction is pixel_z/(grid_z-1) = 1/999
 
-           """
+        """
 
-        xm = torch.linspace(0, pixel_x, grid_x).view(1, 1, 1, -1).expand(1, grid_z, grid_y, grid_x)
-        ym = torch.linspace(0, pixel_y, grid_y).view(1, 1, -1, 1).expand(1, grid_z, grid_y, grid_x)
-        zm = torch.linspace(0, pixel_z, grid_z).view(1, -1, 1, 1).expand(1, grid_z, grid_y, grid_x)
+        xm = (
+            torch.linspace(0, pixel_x, grid_x)
+            .view(1, 1, 1, -1)
+            .expand(1, grid_z, grid_y, grid_x)
+        )
+        ym = (
+            torch.linspace(0, pixel_y, grid_y)
+            .view(1, 1, -1, 1)
+            .expand(1, grid_z, grid_y, grid_x)
+        )
+        zm = (
+            torch.linspace(0, pixel_z, grid_z)
+            .view(1, -1, 1, 1)
+            .expand(1, grid_z, grid_y, grid_x)
+        )
         xyzm = torch.cat((xm, ym, zm), 0)
 
-        self.xyzm = xyzm.cuda()
+        self.xyzm = xyzm.to(device)
         self.one_hot = one_hot
         self.grid_x = grid_x
         self.grid_y = grid_y
@@ -193,8 +254,14 @@ def __init__(self, grid_z, grid_y, grid_x, pixel_z, pixel_y, pixel_x, one_hot=Fa
         self.pixel_x = pixel_x
         self.pixel_y = pixel_y
         self.pixel_z = pixel_z
-
-    def cluster_with_gt(self, prediction, gt_instance, n_sigma=3, ):
+        self.device = device
+
+    def cluster_with_gt(
+        self,
+        prediction,
+        gt_instance,
+        n_sigma=3,
+    ):
         """
         Parameters
         ----------
@@ -209,93 +276,133 @@ def cluster_with_gt(self, prediction, gt_instance, n_sigma=3, ):
         ----------
         instance: PyTorch Tensor (D, H, W)
                 instance segmentation
-       """
+        """
 
-        depth, height, width = prediction.size(1), prediction.size(2), prediction.size(3)
+        depth, height, width = (
+            prediction.size(1),
+            prediction.size(2),
+            prediction.size(3),
+        )
         xyzm_s = self.xyzm[:, 0:depth, 0:height, 0:width]  # 3 x d x h x w
         spatial_emb = torch.tanh(prediction[0:3]) + xyzm_s  # 3 x d x h x w
-        sigma = prediction[3:3 + n_sigma]  # n_sigma x h x w
+        sigma = prediction[3 : 3 + n_sigma]  # n_sigma x h x w
 
-        instance_map = torch.zeros(depth, height, width).short().cuda()
+        instance_map = torch.zeros(depth, height, width).short().to(self.device)
         unique_instances = gt_instance.unique()
         unique_instances = unique_instances[unique_instances != 0]
 
         for id in unique_instances:
             mask = gt_instance.eq(id).view(1, depth, height, width)
-            center = spatial_emb[mask.expand_as(spatial_emb)].view(
-                3, -1).mean(1).view(3, 1, 1, 1)  # 3 x 1 x 1 x 1
-            s = sigma[mask.expand_as(sigma)].view(n_sigma, -1).mean(1).view(n_sigma, 1, 1, 1)
+            center = (
+                spatial_emb[mask.expand_as(spatial_emb)]
+                .view(3, -1)
+                .mean(1)
+                .view(3, 1, 1, 1)
+            )  # 3 x 1 x 1 x 1
+            s = (
+                sigma[mask.expand_as(sigma)]
+                .view(n_sigma, -1)
+                .mean(1)
+                .view(n_sigma, 1, 1, 1)
+            )
 
             s = torch.exp(s * 10)  # n_sigma x 1 x 1
             dist = torch.exp(-1 * torch.sum(torch.pow(spatial_emb - center, 2) * s, 0))
-            proposal = (dist > 0.5)
+            proposal = dist > 0.5
             instance_map[proposal] = id.item()  # TODO
 
         return instance_map
 
-    def cluster(self, prediction, n_sigma=3, seed_thresh=0.9, fg_thresh=0.5, min_mask_sum=0, min_unclustered_sum=0,
-                min_object_size=36):
+    def cluster(
+        self,
+        prediction,
+        n_sigma=3,
+        seed_thresh=0.9,
+        fg_thresh=0.5,
+        min_mask_sum=0,
+        min_unclustered_sum=0,
+        min_object_size=36,
+    ):
+        """
+        Parameters
+        ----------
+        prediction :  PyTorch Tensor
+                Model Prediction (7, D, H, W)
+        n_sigma: int, default = 3
+                Number of dimensions in Raw Image
+        seed_thresh : float, default=0.9
+                Seediness Threshold defines which pixels are considered to
+                identify object centres
+        fg_thresh: float, default=0.5
+                Foreground Threshold defines which pixels are considered to
+                form the "Foreground"
+                and hence which  would need to be clustered into unique objects
+        min_mask_sum: int
+                Only start creating instances, if there are at least
+                `min_mask_sum` pixels in foreground!
+        min_unclustered_sum: int
+                Stop when the number of seed candidates are less than
+                `min_unclustered_sum`
+        min_object_size: int
+            Predicted Objects below this threshold are ignored
 
+        Returns
+        ----------
+        instance: PyTorch Tensor (D, H, W)
+                instance segmentation
         """
-                Parameters
-                ----------
-                prediction :  PyTorch Tensor
-                        Model Prediction (7, D, H, W)
-                n_sigma: int, default = 3
-                        Number of dimensions in Raw Image
-                seed_thresh : float, default=0.9
-                        Seediness Threshold defines which pixels are considered to identify object centres
-                fg_thresh: float, default=0.5
-                        Foreground Threshold defines which pixels are considered to the form the Foreground
-                        and which would need to be clustered into unique objects
-                min_mask_sum: int
-                        Only start creating instances, if there are at least `min_mask_sum` pixels in foreground!
-                min_unclustered_sum: int
-                        Stop when the number of seed candidates are less than `min_unclustered_sum`
-                min_object_size: int
-                    Predicted Objects below this threshold are ignored
-
-                Returns
-                ----------
-                instance: PyTorch Tensor (D, H, W)
-                        instance segmentation
-               """
-
-        depth, height, width = prediction.size(1), prediction.size(2), prediction.size(3)
+
+        depth, height, width = (
+            prediction.size(1),
+            prediction.size(2),
+            prediction.size(3),
+        )
         xyzm_s = self.xyzm[:, 0:depth, 0:height, 0:width]
 
         spatial_emb = torch.tanh(prediction[0:3]) + xyzm_s  # 3 x d x h x w
 
-        sigma = prediction[3:3 + n_sigma]  # n_sigma x d x h x w
-        seed_map = torch.sigmoid(prediction[3 + n_sigma:3 + n_sigma + 1])  # 1 x d x h x w
+        sigma = prediction[3 : 3 + n_sigma]  # n_sigma x d x h x w
+        seed_map = torch.sigmoid(
+            prediction[3 + n_sigma : 3 + n_sigma + 1]
+        )  # 1 x d x h x w
         instance_map = torch.zeros(depth, height, width).short()
 
         count = 1
         mask = seed_map > fg_thresh
-        if mask.sum() > min_mask_sum:  # top level decision: only start creating instances, if there are atleast `min_mask_sum` pixels in foreground!
-
-            spatial_emb_masked = spatial_emb[mask.expand_as(spatial_emb)].view(n_sigma, -1)
+        if (
+            mask.sum() > min_mask_sum
+        ):  # top level decision: only start creating instances,
+            # if there are atleast `min_mask_sum` pixels in foreground!
+            spatial_emb_masked = spatial_emb[mask.expand_as(spatial_emb)].view(
+                n_sigma, -1
+            )
             sigma_masked = sigma[mask.expand_as(sigma)].view(n_sigma, -1)
             seed_map_masked = seed_map[mask].view(1, -1)
 
-            unclustered = torch.ones(mask.sum()).short().cuda()
-            instance_map_masked = torch.zeros(mask.sum()).short().cuda()
+            unclustered = torch.ones(mask.sum()).short().to(self.device)
+            instance_map_masked = torch.zeros(mask.sum()).short().to(self.device)
 
             while (
-                    unclustered.sum() > min_unclustered_sum):  # stop when the seed candidates are less than min_unclustered_sum
+                unclustered.sum() > min_unclustered_sum
+            ):  # stop when the seed candidates are less than min_unclustered_sum
                 seed = (seed_map_masked * unclustered.float()).argmax().item()
                 seed_score = (seed_map_masked * unclustered.float()).max().item()
                 if seed_score < seed_thresh:
                     break
-                center = spatial_emb_masked[:, seed:seed + 1]
+                center = spatial_emb_masked[:, seed : seed + 1]
                 unclustered[seed] = 0
 
-                s = torch.exp(sigma_masked[:, seed:seed + 1] * 10)
-                dist = torch.exp(-1 * torch.sum(torch.pow(spatial_emb_masked - center, 2) * s, 0))
+                s = torch.exp(sigma_masked[:, seed : seed + 1] * 10)
+                dist = torch.exp(
+                    -1 * torch.sum(torch.pow(spatial_emb_masked - center, 2) * s, 0)
+                )
 
                 proposal = (dist > 0.5).squeeze()
                 if proposal.sum() > min_object_size:
-                    if unclustered[proposal].sum().float() / proposal.sum().float() > 0.5:
+                    if (
+                        unclustered[proposal].sum().float() / proposal.sum().float()
+                        > 0.5
+                    ):
                         instance_map_masked[proposal.squeeze()] = count
                         instance_mask = torch.zeros(depth, height, width).short()
                         instance_mask[mask.squeeze().cpu()] = proposal.short().cpu()
@@ -306,38 +413,52 @@ def cluster(self, prediction, n_sigma=3, seed_thresh=0.9, fg_thresh=0.5, min_mas
 
         return instance_map
 
-    def cluster_local_maxima(self, prediction, n_sigma=3, fg_thresh=0.5, min_mask_sum=0, min_unclustered_sum=0,
-                             min_object_size=36):
+    def cluster_local_maxima(
+        self,
+        prediction,
+        n_sigma=3,
+        fg_thresh=0.5,
+        min_mask_sum=0,
+        min_unclustered_sum=0,
+        min_object_size=36,
+    ):
+        """
+        Parameters
+        ----------
+        prediction :  PyTorch Tensor
+                Model Prediction (7, D, H, W)
+        n_sigma: int, default = 3
+                Number of dimensions in Raw Image
+        fg_thresh: float, default=0.5
+                Foreground Threshold defines which pixels are considered to
+                form the "Foreground"
+                and which would need to be clustered into unique objects
+        min_mask_sum: int
+                Only start creating instances, if there are at least
+                `min_mask_sum` pixels in foreground!
+        min_unclustered_sum: int
+                Stop when the number of seed candidates are less
+                than `min_unclustered_sum`
+        min_object_size: int
+            Predicted Objects below this threshold are ignored
 
+        Returns
+        ----------
+        instance: PyTorch Tensor (D, H, W)
+                instance segmentation
         """
-            Parameters
-            ----------
-            prediction :  PyTorch Tensor
-                    Model Prediction (7, D, H, W)
-            n_sigma: int, default = 3
-                    Number of dimensions in Raw Image
-            fg_thresh: float, default=0.5
-                    Foreground Threshold defines which pixels are considered to the form the Foreground
-                    and which would need to be clustered into unique objects
-            min_mask_sum: int
-                    Only start creating instances, if there are at least `min_mask_sum` pixels in foreground!
-            min_unclustered_sum: int
-                    Stop when the number of seed candidates are less than `min_unclustered_sum`
-            min_object_size: int
-                Predicted Objects below this threshold are ignored
-
-            Returns
-            ----------
-            instance: PyTorch Tensor (D, H, W)
-                    instance segmentation
-           """
 
         from scipy.ndimage import gaussian_filter
-        depth, height, width = prediction.size(1), prediction.size(2), prediction.size(3)
+
+        depth, height, width = (
+            prediction.size(1),
+            prediction.size(2),
+            prediction.size(3),
+        )
         xyzm_s = self.xyzm[:, 0:depth, 0:height, 0:width]
         spatial_emb = torch.tanh(prediction[0:3]) + xyzm_s  # 3 x d x h x w
-        sigma = prediction[3:3 + n_sigma]  # n_sigma x d x h x w
-        seed_map = torch.sigmoid(prediction[3 + n_sigma:3 + n_sigma + 1])  # 1 x h x w
+        sigma = prediction[3 : 3 + n_sigma]  # n_sigma x d x h x w
+        seed_map = torch.sigmoid(prediction[3 + n_sigma : 3 + n_sigma + 1])  # 1 x h x w
         instance_map = torch.zeros(depth, height, width).short()
         # instances = []  # list
         count = 1
@@ -351,198 +472,243 @@ def cluster_local_maxima(self, prediction, n_sigma=3, fg_thresh=0.5, min_mask_su
 
         mask_local_max_cpu = np.zeros(seed_map_cpu.shape, dtype=np.bool)
         mask_local_max_cpu[tuple(coords.T)] = True
-        mask_local_max = torch.from_numpy(mask_local_max_cpu).bool().cuda()
+        mask_local_max = torch.from_numpy(mask_local_max_cpu).bool().to(self.device)
 
         mask_seed = mask_fg * mask_local_max
         if mask_fg.sum() > min_mask_sum:
-            spatial_emb_fg_masked = spatial_emb[mask_fg.expand_as(spatial_emb)].view(n_sigma, -1)  # fg candidate pixels
-            spatial_emb_seed_masked = spatial_emb[mask_seed.expand_as(spatial_emb)].view(n_sigma,
-                                                                                         -1)  # seed candidate pixels
-
-            sigma_seed_masked = sigma[mask_seed.expand_as(sigma)].view(n_sigma, -1)  # sigma for seed candidate pixels
-            seed_map_seed_masked = seed_map[mask_seed].view(1, -1)  # seediness for seed candidate pixels
-
-            unprocessed = torch.ones(mask_seed.sum()).short().cuda()  # unprocessed seed candidate pixels
-            unclustered = torch.ones(mask_fg.sum()).short().cuda()  # unclustered fg candidate pixels
-            instance_map_masked = torch.zeros(mask_fg.sum()).short().cuda()
-            while (unprocessed.sum() > min_unclustered_sum):
+            spatial_emb_fg_masked = spatial_emb[mask_fg.expand_as(spatial_emb)].view(
+                n_sigma, -1
+            )  # fg candidate pixels
+            spatial_emb_seed_masked = spatial_emb[
+                mask_seed.expand_as(spatial_emb)
+            ].view(
+                n_sigma, -1
+            )  # seed candidate pixels
+
+            sigma_seed_masked = sigma[mask_seed.expand_as(sigma)].view(
+                n_sigma, -1
+            )  # sigma for seed candidate pixels
+            seed_map_seed_masked = seed_map[mask_seed].view(
+                1, -1
+            )  # seediness for seed candidate pixels
+
+            unprocessed = (
+                torch.ones(mask_seed.sum()).short().to(self.device)
+            )  # unprocessed seed candidate pixels
+            unclustered = (
+                torch.ones(mask_fg.sum()).short().to(self.device)
+            )  # unclustered fg candidate pixels
+            instance_map_masked = torch.zeros(mask_fg.sum()).short().to(self.device)
+            while unprocessed.sum() > min_unclustered_sum:
                 seed = (seed_map_seed_masked * unprocessed.float()).argmax().item()
-                center = spatial_emb_seed_masked[:, seed:seed + 1]
+                center = spatial_emb_seed_masked[:, seed : seed + 1]
                 unprocessed[seed] = 0
-                s = torch.exp(sigma_seed_masked[:, seed:seed + 1] * 10)
-                dist = torch.exp(-1 * torch.sum(torch.pow(spatial_emb_fg_masked - center, 2) * s, 0))
+                s = torch.exp(sigma_seed_masked[:, seed : seed + 1] * 10)
+                dist = torch.exp(
+                    -1 * torch.sum(torch.pow(spatial_emb_fg_masked - center, 2) * s, 0)
+                )
                 proposal = (dist > 0.5).squeeze()
                 if proposal.sum() > min_object_size:
-                    if unclustered[proposal].sum().float() / proposal.sum().float() > 0.5:
+                    if (
+                        unclustered[proposal].sum().float() / proposal.sum().float()
+                        > 0.5
+                    ):
                         instance_map_masked[proposal.squeeze()] = count
                         count += 1
                         unclustered[proposal] = 0
-                        # note the line above increases false positives, tab back twice to show less objects!
-                        # The reason I leave it like so is because the penalty on false-negative nodes is `10` while
-                        # penalty on false-positive nodes is `1`.
             instance_map[mask_fg.squeeze().cpu()] = instance_map_masked.cpu()
         return instance_map
 
 
 class Cluster:
     """
-            A class used to cluster pixel embeddings in 2D
-
-            Attributes
-            ----------
-            xym :  float (2, W, D)
-                    pixel coordinates of tile /grid
-
-            one_hot : bool
-                    Should be set to True, if the GT label masks are present in a one-hot encoded fashion
-
-            grid_x: int
-                    Length (width) of grid / tile
-
-            grid_y: int
-                    Height of grid / tile
-
-            pixel_x: float
-                    if grid_x = 1000 and pixel_x = 1.0, then the pixel spacing along the x direction is pixel_x/(grid_x-1) = 1/999
-            pixel_y: float
-                    if grid_y = 1000 and pixel_y = 1.0, then the pixel spacing along the y direction is pixel_y/(grid_y-1) = 1/999
-
-
-            Methods
-            -------
-            __init__: Initializes an object of class `Cluster_3d`
-
-            cluster_with_gt: use the predicted spatial embeddings from all pixels belonging to the GT label mask
-                        to identify the predicted cluster (used during training and validation)
-
-            cluster:    use the  predicted spatial embeddings from all pixels in the test image.
-                        Employs `fg_thresh` and `seed_thresh`
-            cluster_local_maxima: use the  predicted spatial embeddings from all pixels in the test image.
-                        Employs only `fg_thresh`
-            """
-
-    def __init__(self, grid_y, grid_x, pixel_y, pixel_x, one_hot=False):
+    A class used to cluster pixel embeddings in 2D
+
+    Attributes
+    ----------
+    xym :  float (2, W, D)
+            pixel coordinates of tile /grid
+
+    one_hot : bool
+            Should be set to True, if the GT label masks are present in a
+            one-hot encoded fashion
+
+    grid_x: int
+            Length (width) of grid / tile
+
+    grid_y: int
+            Height of grid / tile
+
+    pixel_x: float
+            if grid_x = 1000 and pixel_x = 1.0, then the pixel spacing
+            along the x direction is pixel_x/(grid_x-1) = 1/999
+    pixel_y: float
+            if grid_y = 1000 and pixel_y = 1.0, then the pixel spacing
+            along the y direction is pixel_y/(grid_y-1) = 1/999
+
+
+    Methods
+    -------
+    __init__: Initializes an object of class `Cluster_3d`
+
+    cluster_with_gt: use the predicted spatial embeddings
+                from all pixels belonging to the GT label mask
+                to identify the predicted cluster
+                (used during training and validation)
+
+    cluster:    use the  predicted spatial embeddings
+                from all pixels in the test image.
+                Employs `fg_thresh` and `seed_thresh`
+    cluster_local_maxima: use the  predicted spatial embeddings
+                from all pixels in the test image.
+                Employs only `fg_thresh`
+    """
 
+    def __init__(self, grid_y, grid_x, pixel_y, pixel_x, device, one_hot=False):
         """
-           Parameters
-           ----------
-           xym :  float (2, W, D)
-                    pixel coordinates of tile /grid
+        Parameters
+        ----------
+        xym :  float (2, W, D)
+                 pixel coordinates of tile /grid
 
-            one_hot : bool
-                    Should be set to True, if the GT label masks are present in a one-hot encoded fashion
+         one_hot : bool
+                 Should be set to True, if the GT label masks are present
+                 in a one-hot encoded fashion
 
-            grid_x: int
-                    Length (width) of grid / tile
+         grid_x: int
+                 Length (width) of grid / tile
 
-            grid_y: int
-                    Height of grid / tile
+         grid_y: int
+                 Height of grid / tile
 
-            pixel_x: float
-                    if grid_x = 1000 and pixel_x = 1.0, then the pixel spacing along the x direction is pixel_x/(grid_x-1) = 1/999
-            pixel_y: float
-                    if grid_y = 1000 and pixel_y = 1.0, then the pixel spacing along the y direction is pixel_y/(grid_y-1) = 1/999
+         pixel_x: float
+                 if grid_x = 1000 and pixel_x = 1.0, then the pixel spacing
+                 along the x direction is pixel_x/(grid_x-1) = 1/999
+         pixel_y: float
+                 if grid_y = 1000 and pixel_y = 1.0, then the pixel spacing
+                 along the y direction is pixel_y/(grid_y-1) = 1/999
 
-                   """
+        """
         xm = torch.linspace(0, pixel_x, grid_x).view(1, 1, -1).expand(1, grid_y, grid_x)
         ym = torch.linspace(0, pixel_y, grid_y).view(1, -1, 1).expand(1, grid_y, grid_x)
         xym = torch.cat((xm, ym), 0)
 
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        if device == "cuda":
-            self.xym = xym.cuda()
-        else:
-            self.xym = xym
+        self.device = device
+        self.xym = xym.to(self.device)
         self.one_hot = one_hot
         self.grid_x = grid_x
         self.grid_y = grid_y
         self.pixel_x = pixel_x
         self.pixel_y = pixel_y
 
-    def cluster_with_gt(self, prediction, gt_instance, n_sigma=2, ):
+    def cluster_with_gt(
+        self,
+        prediction,
+        gt_instance,
+        n_sigma=2,
+    ):
+        """
+        Parameters
+        ----------
+        prediction :  PyTorch Tensor
+                Model Prediction (5, H, W)
+        gt_instance : PyTorch Tensor
+                Ground Truth Instance Segmentation Label Map
+
+        n_sigma: int, default = 2
+                Number of dimensions in Raw Image
+        Returns
+        ----------
+        instance: PyTorch Tensor (H, W)
+                instance segmentation
         """
-            Parameters
-            ----------
-            prediction :  PyTorch Tensor
-                    Model Prediction (5, H, W)
-            gt_instance : PyTorch Tensor
-                    Ground Truth Instance Segmentation Label Map
-
-            n_sigma: int, default = 2
-                    Number of dimensions in Raw Image
-            Returns
-            ----------
-            instance: PyTorch Tensor (H, W)
-                    instance segmentation
-       """
 
         height, width = prediction.size(1), prediction.size(2)
 
         xym_s = self.xym[:, 0:height, 0:width]  # 2 x h x w
 
         spatial_emb = torch.tanh(prediction[0:2]) + xym_s  # 2 x h x w
-        sigma = prediction[2:2 + n_sigma]  # n_sigma x h x w
+        sigma = prediction[2 : 2 + n_sigma]  # n_sigma x h x w
 
-        instance_map = torch.zeros(height, width).short().cuda()
-        if (self.one_hot):
+        instance_map = torch.zeros(height, width).short().to(self.device)
+        if self.one_hot:
             unique_instances = torch.arange(gt_instance.size(0))
         else:
             unique_instances = gt_instance.unique()
             unique_instances = unique_instances[unique_instances != 0]
 
         for id in unique_instances:
-            if (self.one_hot):
+            if self.one_hot:
                 mask = gt_instance[id].eq(1).view(1, height, width)
             else:
                 mask = gt_instance.eq(id).view(1, height, width)
 
-            center = spatial_emb[mask.expand_as(spatial_emb)].view(
-                2, -1).mean(1).view(2, 1, 1)  # 2 x 1 x 1
+            center = (
+                spatial_emb[mask.expand_as(spatial_emb)]
+                .view(2, -1)
+                .mean(1)
+                .view(2, 1, 1)
+            )  # 2 x 1 x 1
 
-            s = sigma[mask.expand_as(sigma)].view(n_sigma, -1).mean(1).view(n_sigma, 1, 1)
+            s = (
+                sigma[mask.expand_as(sigma)]
+                .view(n_sigma, -1)
+                .mean(1)
+                .view(n_sigma, 1, 1)
+            )
 
             s = torch.exp(s * 10)  # n_sigma x 1 x 1 #
             dist = torch.exp(-1 * torch.sum(torch.pow(spatial_emb - center, 2) * s, 0))
-            proposal = (dist > 0.5)
-            if (self.one_hot):
+            proposal = dist > 0.5
+            if self.one_hot:
                 instance_map[proposal] = id.item() + 1
             else:
                 instance_map[proposal] = id.item()
 
         return instance_map
 
-    def cluster_local_maxima(self, prediction, n_sigma=2, fg_thresh=0.5, min_mask_sum=0, min_unclustered_sum=0,
-                             min_object_size=36):
+    def cluster_local_maxima(
+        self,
+        prediction,
+        n_sigma=2,
+        fg_thresh=0.5,
+        min_mask_sum=0,
+        min_unclustered_sum=0,
+        min_object_size=36,
+    ):
+        """
+        Parameters
+        ----------
+        prediction :  PyTorch Tensor
+                Model Prediction (5, H, W)
+        n_sigma: int, default = 2
+                Number of dimensions in Raw Image
+        fg_thresh: float, default=0.5
+                Foreground Threshold defines which pixels are considered
+                to form the Foreground
+                and which would need to be clustered into unique objects
+        min_mask_sum: int
+                Only start creating instances, if there are at least
+                `min_mask_sum` pixels in foreground!
+        min_unclustered_sum: int
+                Stop when the number of seed candidates are less
+                than `min_unclustered_sum`
+        min_object_size: int
+            Predicted Objects below this threshold are ignored
 
+        Returns
+        ----------
+        instance: PyTorch Tensor (H, W)
+                instance segmentation
         """
-            Parameters
-            ----------
-            prediction :  PyTorch Tensor
-                    Model Prediction (5, H, W)
-            n_sigma: int, default = 2
-                    Number of dimensions in Raw Image
-            fg_thresh: float, default=0.5
-                    Foreground Threshold defines which pixels are considered to the form the Foreground
-                    and which would need to be clustered into unique objects
-            min_mask_sum: int
-                    Only start creating instances, if there are at least `min_mask_sum` pixels in foreground!
-            min_unclustered_sum: int
-                    Stop when the number of seed candidates are less than `min_unclustered_sum`
-            min_object_size: int
-                Predicted Objects below this threshold are ignored
-
-            Returns
-            ----------
-            instance: PyTorch Tensor (H, W)
-                    instance segmentation
-       """
 
         from scipy.ndimage import gaussian_filter
+
         height, width = prediction.size(1), prediction.size(2)
         xym_s = self.xym[:, 0:height, 0:width]
         spatial_emb = torch.tanh(prediction[0:2]) + xym_s  # 2 x h x w
-        sigma = prediction[2:2 + n_sigma]  # n_sigma x h x w
-        seed_map = torch.sigmoid(prediction[2 + n_sigma:2 + n_sigma + 1])  # 1 x h x w
+        sigma = prediction[2 : 2 + n_sigma]  # n_sigma x h x w
+        seed_map = torch.sigmoid(prediction[2 + n_sigma : 2 + n_sigma + 1])  # 1 x h x w
         instance_map = torch.zeros(height, width).short()
         # instances = []  # list
         count = 1
@@ -555,73 +721,100 @@ def cluster_local_maxima(self, prediction, n_sigma=2, fg_thresh=0.5, min_mask_su
 
         mask_local_max_cpu = np.zeros(seed_map_cpu.shape, dtype=np.bool)
         mask_local_max_cpu[tuple(coords.T)] = True
-        mask_local_max = torch.from_numpy(mask_local_max_cpu).bool().cuda()
+        mask_local_max = torch.from_numpy(mask_local_max_cpu).bool().to(self.device)
 
         mask_seed = mask_fg * mask_local_max
         if mask_fg.sum() > min_mask_sum:
-            spatial_emb_fg_masked = spatial_emb[mask_fg.expand_as(spatial_emb)].view(n_sigma, -1)  # fg candidate pixels
-            spatial_emb_seed_masked = spatial_emb[mask_seed.expand_as(spatial_emb)].view(n_sigma,
-                                                                                         -1)  # seed candidate pixels
-
-            sigma_seed_masked = sigma[mask_seed.expand_as(sigma)].view(n_sigma, -1)  # sigma for seed candidate pixels
-            seed_map_seed_masked = seed_map[mask_seed].view(1, -1)  # seediness for seed candidate pixels
-
-            unprocessed = torch.ones(mask_seed.sum()).short().cuda()  # unclustered seed candidate pixels
-            unclustered = torch.ones(mask_fg.sum()).short().cuda()  # unclustered fg candidate pixels
-            instance_map_masked = torch.zeros(mask_fg.sum()).short().cuda()
-            while (unprocessed.sum() > min_unclustered_sum):
+            spatial_emb_fg_masked = spatial_emb[mask_fg.expand_as(spatial_emb)].view(
+                n_sigma, -1
+            )  # fg candidate pixels
+            spatial_emb_seed_masked = spatial_emb[
+                mask_seed.expand_as(spatial_emb)
+            ].view(
+                n_sigma, -1
+            )  # seed candidate pixels
+
+            sigma_seed_masked = sigma[mask_seed.expand_as(sigma)].view(
+                n_sigma, -1
+            )  # sigma for seed candidate pixels
+            seed_map_seed_masked = seed_map[mask_seed].view(
+                1, -1
+            )  # seediness for seed candidate pixels
+
+            unprocessed = (
+                torch.ones(mask_seed.sum()).short().to(self.device)
+            )  # unclustered seed candidate pixels
+            unclustered = (
+                torch.ones(mask_fg.sum()).short().to(self.device)
+            )  # unclustered fg candidate pixels
+            instance_map_masked = torch.zeros(mask_fg.sum()).short().to(self.device)
+            while unprocessed.sum() > min_unclustered_sum:
                 seed = (seed_map_seed_masked * unprocessed.float()).argmax().item()
-                center = spatial_emb_seed_masked[:, seed:seed + 1]
+                center = spatial_emb_seed_masked[:, seed : seed + 1]
                 unprocessed[seed] = 0
-                s = torch.exp(sigma_seed_masked[:, seed:seed + 1] * 10)
-                dist = torch.exp(-1 * torch.sum(torch.pow(spatial_emb_fg_masked - center, 2) * s, 0))
+                s = torch.exp(sigma_seed_masked[:, seed : seed + 1] * 10)
+                dist = torch.exp(
+                    -1 * torch.sum(torch.pow(spatial_emb_fg_masked - center, 2) * s, 0)
+                )
                 proposal = (dist > 0.5).squeeze()
                 if proposal.sum() > min_object_size:
-                    if unclustered[proposal].sum().float() / proposal.sum().float() > 0.5:
+                    if (
+                        unclustered[proposal].sum().float() / proposal.sum().float()
+                        > 0.5
+                    ):
                         instance_map_masked[proposal.squeeze()] = count
                         count += 1
                         unclustered[proposal] = 0
-                        # note the line above increases false positives, tab back twice to show less objects!
-                        # The reason I leave it like so is because the penalty on false-negative nodes is `10` while
-                        # penalty on false-positive nodes is `1`.
             instance_map[mask_fg.squeeze().cpu()] = instance_map_masked.cpu()
         return instance_map
 
-    def cluster(self, prediction, n_sigma=2, seed_thresh=0.9, fg_thresh=0.5, min_mask_sum=0, min_unclustered_sum=0,
-                min_object_size=36):
+    def cluster(
+        self,
+        prediction,
+        n_sigma=2,
+        seed_thresh=0.9,
+        fg_thresh=0.5,
+        min_mask_sum=0,
+        min_unclustered_sum=0,
+        min_object_size=36,
+    ):
+        """
+        Parameters
+        ----------
+        prediction :  PyTorch Tensor
+                Model Prediction (5, H, W)
+        n_sigma: int, default = 2
+                Number of dimensions in Raw Image
+        seed_thresh : float, default=0.9
+                Seediness Threshold defines which pixels are considered to
+                identify object centres
+        fg_thresh: float, default=0.5
+                Foreground Threshold defines which pixels are considered to
+                form the "foreground"
+                and hence would need to be clustered into unique objects
+        min_mask_sum: int
+                Only start creating instances, if there are at least
+                `min_mask_sum` pixels in foreground!
+        min_unclustered_sum: int
+                Stop when the number of seed candidates are less than
+                `min_unclustered_sum`
+        min_object_size: int
+            Predicted Objects below this threshold are ignored
 
+        Returns
+        ----------
+        instance: PyTorch Tensor (H, W)
+                instance segmentation
         """
-            Parameters
-            ----------
-            prediction :  PyTorch Tensor
-                    Model Prediction (5, H, W)
-            n_sigma: int, default = 2
-                    Number of dimensions in Raw Image
-            seed_thresh : float, default=0.9
-                    Seediness Threshold defines which pixels are considered to identify object centres
-            fg_thresh: float, default=0.5
-                    Foreground Threshold defines which pixels are considered to the form the Foreground
-                    and which would need to be clustered into unique objects
-            min_mask_sum: int
-                    Only start creating instances, if there are at least `min_mask_sum` pixels in foreground!
-            min_unclustered_sum: int
-                    Stop when the number of seed candidates are less than `min_unclustered_sum`
-            min_object_size: int
-                Predicted Objects below this threshold are ignored
-
-            Returns
-            ----------
-            instance: PyTorch Tensor (H, W)
-                    instance segmentation
-           """
 
         height, width = prediction.size(1), prediction.size(2)
         xym_s = self.xym[:, 0:height, 0:width]
 
         spatial_emb = torch.tanh(prediction[0:2]) + xym_s  # 2 x h x w
 
-        sigma = prediction[2:2 + n_sigma]  # n_sigma x h x w
-        seed_map = torch.sigmoid(prediction[2 + n_sigma:2 + n_sigma + 1])  # 1 x h x w
+        sigma = prediction[2 : 2 + n_sigma]  # n_sigma x h x w
+        seed_map = torch.sigmoid(prediction[2 + n_sigma : 2 + n_sigma + 1])
+        # 1 x h x w
 
         instance_map = torch.zeros(height, width).short()
         # instances = []  # list
@@ -630,40 +823,37 @@ def cluster(self, prediction, n_sigma=2, seed_thresh=0.9, fg_thresh=0.5, min_mas
         mask = seed_map > fg_thresh
 
         if mask.sum() > min_mask_sum:
-
-            spatial_emb_masked = spatial_emb[mask.expand_as(spatial_emb)].view(n_sigma, -1)
+            spatial_emb_masked = spatial_emb[mask.expand_as(spatial_emb)].view(
+                n_sigma, -1
+            )
             sigma_masked = sigma[mask.expand_as(sigma)].view(n_sigma, -1)
             seed_map_masked = seed_map[mask].view(1, -1)
 
-            unclustered = torch.ones(mask.sum()).short().cuda()
-            instance_map_masked = torch.zeros(mask.sum()).short().cuda()
+            unclustered = torch.ones(mask.sum()).short().to(self.device)
+            instance_map_masked = torch.zeros(mask.sum()).short().to(self.device)
 
-            while (unclustered.sum() > min_unclustered_sum):
+            while unclustered.sum() > min_unclustered_sum:
                 seed = (seed_map_masked * unclustered.float()).argmax().item()
                 seed_score = (seed_map_masked * unclustered.float()).max().item()
                 if seed_score < seed_thresh:
                     break
-                center = spatial_emb_masked[:, seed:seed + 1]
+                center = spatial_emb_masked[:, seed : seed + 1]
                 unclustered[seed] = 0
 
-                s = torch.exp(sigma_masked[:, seed:seed + 1] * 10)
-                dist = torch.exp(-1 * torch.sum(torch.pow(spatial_emb_masked - center, 2) * s, 0))
+                s = torch.exp(sigma_masked[:, seed : seed + 1] * 10)
+                dist = torch.exp(
+                    -1 * torch.sum(torch.pow(spatial_emb_masked - center, 2) * s, 0)
+                )
 
                 proposal = (dist > 0.5).squeeze()
                 if proposal.sum() > min_object_size:
-                    if unclustered[proposal].sum().float() / proposal.sum().float() > 0.5:
+                    if (
+                        unclustered[proposal].sum().float() / proposal.sum().float()
+                        > 0.5
+                    ):
                         instance_map_masked[proposal.squeeze()] = count
                         instance_mask = torch.zeros(height, width).short()
                         instance_mask[mask.squeeze().cpu()] = proposal.short().cpu()
-                        # center_image = torch.zeros(height, width).short()
-                        #
-                        # center[0] = int(degrid(center[0].cpu().detach().numpy(), self.grid_x, self.pixel_x))
-                        # center[1] = int(degrid(center[1].cpu().detach().numpy(), self.grid_y, self.pixel_y))
-                        # center_image[np.clip(int(center[1].item()), 0, height - 1), np.clip(int(center[0].item()), 0,
-                        #                                                                     width - 1)] = True
-                        # instances.append(
-                        #     {'mask': instance_mask.squeeze() * 255, 'score': seed_score,
-                        #      'center-image': center_image})
                         count += 1
 
                 unclustered[proposal] = 0
@@ -674,17 +864,14 @@ def cluster(self, prediction, n_sigma=2, seed_thresh=0.9, fg_thresh=0.5, min_mas
 
 
 class Logger:
-
     def __init__(self, keys, title=""):
-
         self.data = {k: [] for k in keys}
         self.title = title
         self.win = None
 
-        print('Created logger with keys:  {}'.format(keys))
+        print("Created logger with keys:  {}".format(keys))
 
     def plot(self, save=False, save_dir=""):
-
         if self.win is None:
             self.win = plt.subplots()
         fig, ax = self.win
@@ -693,12 +880,12 @@ def plot(self, save=False, save_dir=""):
         keys = []
         count = 0
         for key in self.data:
-            if (count < 3):
+            if count < 3:
                 keys.append(key)
                 data = self.data[key]
-                ax.plot(range(len(data)), data, marker='.')
+                ax.plot(range(len(data)), data, marker=".")
                 count += 1
-        ax.legend(keys, loc='upper right')
+        ax.legend(keys, loc="upper right")
         ax.set_title(self.title)
 
         plt.draw()
@@ -707,11 +894,11 @@ def plot(self, save=False, save_dir=""):
 
         if save:
             # save figure
-            fig.savefig(os.path.join(save_dir, self.title + '.png'))
+            fig.savefig(os.path.join(save_dir, self.title + ".png"))
 
             # save data as csv
             df = pd.DataFrame.from_dict(self.data)
-            df.to_csv(os.path.join(save_dir, self.title + '.csv'))
+            df.to_csv(os.path.join(save_dir, self.title + ".csv"))
 
     def add(self, key, value):
         assert key in self.data, "Key not in data"
@@ -729,16 +916,26 @@ def add_samples(samples, ax, n, amax):
     return samples_list
 
 
-def prepare_embedding_for_train_image(one_hot, grid_x, grid_y, pixel_x, pixel_y, predictions, instance_ids,
-                                      center_images, output,
-                                      instances, n_sigma):
+def prepare_embedding_for_train_image(
+    one_hot,
+    grid_x,
+    grid_y,
+    pixel_x,
+    pixel_y,
+    predictions,
+    instance_ids,
+    center_images,
+    output,
+    instances,
+    n_sigma,
+):
     xm = torch.linspace(0, pixel_x, grid_x).view(1, 1, -1).expand(1, grid_y, grid_x)
     ym = torch.linspace(0, pixel_y, grid_y).view(1, -1, 1).expand(1, grid_y, grid_x)
     xym = torch.cat((xm, ym), 0)
     height, width = predictions.size(0), predictions.size(1)
     xym_s = xym[:, 0:height, 0:width].contiguous()
     spatial_emb = torch.tanh(output[0, 0:2]).cpu() + xym_s
-    sigma = output[0, 2: 2 + n_sigma]  # 2/3 Y X
+    sigma = output[0, 2 : 2 + n_sigma]  # 2/3 Y X
     color_sample = sns.color_palette("dark")
     color_embedding = sns.color_palette("bright")
     color_sample_dic = {}
@@ -755,8 +952,10 @@ def prepare_embedding_for_train_image(one_hot, grid_x, grid_y, pixel_x, pixel_y,
         instance_ids += 1  # make the ids go from 1 to ...
 
     for id in instance_ids:
-        if (one_hot):
-            in_mask = instances[0, id - 1, ...].eq(1)  # for one_hot, id goes from 1 to ...
+        if one_hot:
+            in_mask = instances[0, id - 1, ...].eq(
+                1
+            )  # for one_hot, id goes from 1 to ...
         else:
             in_mask = instances[0].eq(id)  # 1 x h x w
 
@@ -771,14 +970,16 @@ def prepare_embedding_for_train_image(one_hot, grid_x, grid_y, pixel_x, pixel_y,
         spatial_emb_in = spatial_emb[in_mask.expand_as(spatial_emb)].view(2, -1)
         samples_spatial_embeddings = spatial_emb_in[:, idx]
 
-        sample_spatial_embedding_x[id.item()] = add_samples(samples_spatial_embeddings, 0, grid_x - 1,
-                                                            pixel_x)
-        sample_spatial_embedding_y[id.item()] = add_samples(samples_spatial_embeddings, 1, grid_y - 1,
-                                                            pixel_y)
+        sample_spatial_embedding_x[id.item()] = add_samples(
+            samples_spatial_embeddings, 0, grid_x - 1, pixel_x
+        )
+        sample_spatial_embedding_y[id.item()] = add_samples(
+            samples_spatial_embeddings, 1, grid_y - 1, pixel_y
+        )
 
         # center_mask = in_mask & center_images[0].byte()
         center_mask = in_mask & center_images[0]
-        if (center_mask.sum().eq(1)):
+        if center_mask.sum().eq(1):
             center = xym_s[center_mask.expand_as(xym_s)].view(2, 1, 1)
         else:
             xy_in = xym_s[in_mask.expand_as(xym_s)].view(2, -1)
@@ -798,11 +999,23 @@ def prepare_embedding_for_train_image(one_hot, grid_x, grid_y, pixel_x, pixel_y,
         # colors
         color_sample_dic[id.item()] = color_sample[int(id % 10)]
         color_embedding_dic[id.item()] = color_embedding[int(id % 10)]
-    return center_x, center_y, samples_x, samples_y, sample_spatial_embedding_x, \
-           sample_spatial_embedding_y, sigma_x, sigma_y, color_sample_dic, color_embedding_dic
-
-
-def prepare_embedding_for_test_image(instance_map, output, grid_x, grid_y, pixel_x, pixel_y, predictions, n_sigma):
+    return (
+        center_x,
+        center_y,
+        samples_x,
+        samples_y,
+        sample_spatial_embedding_x,
+        sample_spatial_embedding_y,
+        sigma_x,
+        sigma_y,
+        color_sample_dic,
+        color_embedding_dic,
+    )
+
+
+def prepare_embedding_for_test_image(
+    instance_map, output, grid_x, grid_y, pixel_x, pixel_y, predictions, n_sigma
+):
     instance_ids = instance_map.unique()
     instance_ids = instance_ids[instance_ids != 0]
 
@@ -812,7 +1025,7 @@ def prepare_embedding_for_test_image(instance_map, output, grid_x, grid_y, pixel
     height, width = instance_map.size(0), instance_map.size(1)
     xym_s = xym[:, 0:height, 0:width].contiguous()
     spatial_emb = torch.tanh(output[0, 0:2]).cpu() + xym_s
-    sigma = output[0, 2:2 + n_sigma]
+    sigma = output[0, 2 : 2 + n_sigma]
     color_sample = sns.color_palette("dark")
     color_embedding = sns.color_palette("bright")
     color_sample_dic = {}
@@ -838,12 +1051,18 @@ def prepare_embedding_for_test_image(instance_map, output, grid_x, grid_y, pixel
         spatial_emb_in = spatial_emb[in_mask.expand_as(spatial_emb)].view(2, -1)
         samples_spatial_embeddings = spatial_emb_in[:, idx]
 
-        sample_spatial_embedding_x[id.item()] = add_samples(samples_spatial_embeddings, 0, grid_x - 1, pixel_x)
-        sample_spatial_embedding_y[id.item()] = add_samples(samples_spatial_embeddings, 1, grid_y - 1, pixel_y)
-        center_image = predictions[id.item() - 1]['center-image']  # predictions is a list!
+        sample_spatial_embedding_x[id.item()] = add_samples(
+            samples_spatial_embeddings, 0, grid_x - 1, pixel_x
+        )
+        sample_spatial_embedding_y[id.item()] = add_samples(
+            samples_spatial_embeddings, 1, grid_y - 1, pixel_y
+        )
+        center_image = predictions[id.item() - 1][
+            "center-image"
+        ]  # predictions is a list!
         center_mask = in_mask & center_image.byte()
 
-        if (center_mask.sum().eq(1)):
+        if center_mask.sum().eq(1):
             center = xym_s[center_mask.expand_as(xym_s)].view(2, 1, 1)
         else:
             xy_in = xym_s[in_mask.expand_as(xym_s)].view(2, -1)
@@ -864,5 +1083,15 @@ def prepare_embedding_for_test_image(instance_map, output, grid_x, grid_y, pixel
         color_sample_dic[id.item()] = color_sample[int(id % 10)]
         color_embedding_dic[id.item()] = color_embedding[int(id % 10)]
 
-    return center_x, center_y, samples_x, samples_y, sample_spatial_embedding_x, \
-           sample_spatial_embedding_y, sigma_x, sigma_y, color_sample_dic, color_embedding_dic
+    return (
+        center_x,
+        center_y,
+        samples_x,
+        samples_y,
+        sample_spatial_embedding_x,
+        sample_spatial_embedding_y,
+        sigma_x,
+        sigma_y,
+        color_sample_dic,
+        color_embedding_dic,
+    )
diff --git a/EmbedSeg/utils/visualize.py b/EmbedSeg/utils/visualize.py
index e3c3d49..f6dcc5e 100644
--- a/EmbedSeg/utils/visualize.py
+++ b/EmbedSeg/utils/visualize.py
@@ -17,25 +17,27 @@
 
 def create_color_map(n_colors=10):
     """
-        Create Glasbey-style Color map
+    Create Glasbey-style Color map
 
-        Parameters
-        -------
+    Parameters
+    -------
 
-        n_colors: int
+    n_colors: int
 
 
-        Returns
-        -------
-        Color Map
+    Returns
+    -------
+    Color Map
 
-        """
+    """
 
-    gb = Glasbey(base_palette=[(255, 0, 0), (0, 255, 0), (0, 0, 255)],
-                 lightness_range=(10, 100),
-                 hue_range=(10, 100),
-                 chroma_range=(10, 100),
-                 no_black=True)
+    gb = Glasbey(
+        base_palette=[(255, 0, 0), (0, 255, 0), (0, 0, 255)],
+        lightness_range=(10, 100),
+        hue_range=(10, 100),
+        chroma_range=(10, 100),
+        no_black=True,
+    )
     p = gb.generate_palette(size=n_colors)
     p[0, :] = [0, 0, 0]  # make label 0 always black!
     p_ = np.hstack((p, np.ones((p.shape[0], 1))))
@@ -46,53 +48,55 @@ def create_color_map(n_colors=10):
 
 def visualize(image, prediction, ground_truth, seed, new_cmp):
     """
-        Visualizes 2 x 2 grid with Top-Left (Image), Top-Right (Ground Truth), Bottom-Left (Seed),
-        Bottom-Right (Instance Segmentation Prediction)
+    Visualizes 2 x 2 grid with Top-Left (Image),
+    Top-Right (Ground Truth), Bottom-Left (Seed),
+    Bottom-Right (Instance Segmentation Prediction)
 
-        Parameters
-        -------
+    Parameters
+    -------
 
-        image: Numpy Array (YX or 1YX)
-            Raw Image
-        prediction: Numpy Array (YX)
-            Model Prediction of Instance Segmentation
-        ground_truth: Numpy Array (YX)
-            GT Label Mask
-        seed: Confidence (Seediness) Map
+    image: Numpy Array (YX or 1YX)
+        Raw Image
+    prediction: Numpy Array (YX)
+        Model Prediction of Instance Segmentation
+    ground_truth: Numpy Array (YX)
+        GT Label Mask
+    seed: Confidence (Seediness) Map
 
-        new_cmp: Color Map
+    new_cmp: Color Map
 
-        Returns
-        -------
+    Returns
+    -------
 
-        """
+    """
 
-    font = {'family': 'serif',
-            'color': 'white',
-            'weight': 'bold',
-            'size': 16,
-            }
+    font = {
+        "family": "serif",
+        "color": "white",
+        "weight": "bold",
+        "size": 16,
+    }
     plt.figure(figsize=(15, 15))
     img_show = image if image.ndim == 2 else image[0, ...]
-    plt.subplot(221);
-    plt.imshow(img_show, cmap='magma');
+    plt.subplot(221)
+    plt.imshow(img_show, cmap="magma")
     plt.text(30, 30, "IM", fontdict=font)
-    plt.xlabel('Image')
-    plt.axis('off')
-    if (ground_truth is not None):
-        plt.subplot(222);
-        plt.axis('off')
-        plt.imshow(ground_truth, cmap=new_cmp, interpolation='None')
+    plt.xlabel("Image")
+    plt.axis("off")
+    if ground_truth is not None:
+        plt.subplot(222)
+        plt.axis("off")
+        plt.imshow(ground_truth, cmap=new_cmp, interpolation="None")
         plt.text(30, 30, "GT", fontdict=font)
-        plt.xlabel('Ground Truth')
-    plt.subplot(223);
-    plt.axis('off')
-    plt.imshow(seed, interpolation='None')
-    plt.subplot(224);
-    plt.axis('off')
-    plt.imshow(prediction, cmap=new_cmp, interpolation='None')
+        plt.xlabel("Ground Truth")
+    plt.subplot(223)
+    plt.axis("off")
+    plt.imshow(seed, interpolation="None")
+    plt.subplot(224)
+    plt.axis("off")
+    plt.imshow(prediction, cmap=new_cmp, interpolation="None")
     plt.text(30, 30, "PRED", fontdict=font)
-    plt.xlabel('Prediction')
+    plt.xlabel("Prediction")
     plt.tight_layout()
     plt.show()
 
@@ -105,59 +109,75 @@ def decode(filename, one_hot=False, center=False):
     if one_hot:
         mask_decoded = []
         for row in df_numpy:
-            d['counts'] = ast.literal_eval(row[1])
-            d['size'] = ast.literal_eval(row[2])
+            d["counts"] = ast.literal_eval(row[1])
+            d["size"] = ast.literal_eval(row[2])
             mask = rletools.decode(d)  # returns binary mask
             mask_decoded.append(mask)
     else:
         if center:
-            mask_decoded = np.zeros(ast.literal_eval(df_numpy[0][2]),
-                                    dtype=np.bool)  # obtain size by reading first row of csv file
+            mask_decoded = np.zeros(
+                ast.literal_eval(df_numpy[0][2]), dtype=np.bool
+            )  # obtain size by reading first row of csv file
         else:
-            mask_decoded = np.zeros(ast.literal_eval(df_numpy[0][2]),
-                                    dtype=np.uint16)  # obtain size by reading first row of csv file
+            mask_decoded = np.zeros(
+                ast.literal_eval(df_numpy[0][2]), dtype=np.uint16
+            )  # obtain size by reading first row of csv file
         for row in df_numpy:
-            d['counts'] = ast.literal_eval(row[1])
-            d['size'] = ast.literal_eval(row[2])
+            d["counts"] = ast.literal_eval(row[1])
+            d["size"] = ast.literal_eval(row[2])
             mask = rletools.decode(d)  # returns binary mask
             y, x = np.where(mask == 1)
             mask_decoded[y, x] = int(row[0])
     return np.asarray(mask_decoded)
 
 
-def visualize_many_crops(data_dir, project_name, train_val_dir, center, n_images, new_cmp, one_hot=False):
+def visualize_many_crops(
+    data_dir, project_name, train_val_dir, center, n_images, new_cmp, one_hot=False
+):
     """
-            Parameters
-            -------
-            data_dir: str
-                Place where crops are (for example: data_dir='crops')
-            project_name: str
-                Name of Project (for example: project_name= 'dsb-2018')
-            train_val_dir: str
-                One of 'train' or 'val'
-            center: str
-                One of 'medoid', 'centroid', 'approximate-medoid'
-            n_images: int
-                Number of columns
-
-            new_cmp: Color Map
-
-            one_hot: bool
-                If GT label masks are available in one-hot encoded fashion, then this parameter should be set equal to True
-            Returns
-            -------
-
-            """
-
-    font = {'family': 'serif',
-            'color': 'black',
-            'weight': 'bold',
-            'size': 16,
-            }
-    im_filenames = sorted(glob(os.path.join(data_dir, project_name, train_val_dir) + '/images/*.tif'))
-    ma_filenames = sorted(glob(os.path.join(data_dir, project_name, train_val_dir) + '/masks/*'))  # `tifs` or `csvs`
+    Parameters
+    -------
+    data_dir: str
+        Place where crops are (for example: data_dir='crops')
+    project_name: str
+        Name of Project (for example: project_name= 'dsb-2018')
+    train_val_dir: str
+        One of 'train' or 'val'
+    center: str
+        One of 'medoid', 'centroid', 'approximate-medoid'
+    n_images: int
+        Number of columns
+
+    new_cmp: Color Map
+
+    one_hot: bool
+        If GT label masks are available in one-hot encoded fashion,
+        then this parameter should be set equal to True
+    Returns
+    -------
+
+    """
+
+    font = {
+        "family": "serif",
+        "color": "black",
+        "weight": "bold",
+        "size": 16,
+    }
+    im_filenames = sorted(
+        glob(os.path.join(data_dir, project_name, train_val_dir) + "/images/*.tif")
+    )
+    ma_filenames = sorted(
+        glob(os.path.join(data_dir, project_name, train_val_dir) + "/masks/*")
+    )  # `tifs` or `csvs`
     center_filenames = sorted(
-        glob(os.path.join(data_dir, project_name, train_val_dir) + '/center-' + center + '/*'))  # `tifs` or `csvs`
+        glob(
+            os.path.join(data_dir, project_name, train_val_dir)
+            + "/center-"
+            + center
+            + "/*"
+        )
+    )  # `tifs` or `csvs`
     indices = np.random.randint(0, len(im_filenames), n_images)
     fig = plt.figure(constrained_layout=False, figsize=(16, 10))
     spec = gridspec.GridSpec(ncols=n_images, nrows=3, figure=fig)
@@ -165,81 +185,103 @@ def visualize_many_crops(data_dir, project_name, train_val_dir, center, n_images
         ax0 = fig.add_subplot(spec[0, i])
         im = tifffile.imread(im_filenames[index])
         if im.ndim == 2:
-            ax0.imshow(im, cmap='magma', interpolation='None')
+            ax0.imshow(im, cmap="magma", interpolation="None")
         else:
-            ax0.imshow(im[0], cmap='magma', interpolation='None')
+            ax0.imshow(im[0], cmap="magma", interpolation="None")
         ax0.axes.get_xaxis().set_visible(False)
         ax0.set_yticklabels([])
         ax0.set_yticks([])
         if i == 0:
-            ax0.set_ylabel('IM', fontdict=font)
+            ax0.set_ylabel("IM", fontdict=font)
         ax1 = fig.add_subplot(spec[1, i])
         if one_hot:
-            if (ma_filenames[index][-3:] == 'csv'):
+            if ma_filenames[index][-3:] == "csv":
                 label = invert_one_hot(decode(ma_filenames[index], one_hot=True))
             else:
                 label = invert_one_hot(tifffile.imread(ma_filenames[index]))
         else:
-            if (ma_filenames[index][-3:] == 'csv'):
+            if ma_filenames[index][-3:] == "csv":
                 label, _, _ = relabel_sequential(decode(ma_filenames[index]))
             else:
                 label, _, _ = relabel_sequential(tifffile.imread(ma_filenames[index]))
-        ax1.imshow(label, cmap=new_cmp, interpolation='None')
+        ax1.imshow(label, cmap=new_cmp, interpolation="None")
         ax1.axes.get_xaxis().set_visible(False)
         ax1.set_yticklabels([])
         ax1.set_yticks([])
         if i == 0:
-            ax1.set_ylabel('MASK', fontdict=font)
+            ax1.set_ylabel("MASK", fontdict=font)
         ax2 = fig.add_subplot(spec[2, i])
-        if center_filenames[index][-3:] == 'csv':
-            ax2.imshow(decode(center_filenames[index], center=True), cmap='gray', interpolation='None')
+        if center_filenames[index][-3:] == "csv":
+            ax2.imshow(
+                decode(center_filenames[index], center=True),
+                cmap="gray",
+                interpolation="None",
+            )
         else:
-            ax2.imshow(tifffile.imread(center_filenames[index]), cmap='gray', interpolation='None')
+            ax2.imshow(
+                tifffile.imread(center_filenames[index]),
+                cmap="gray",
+                interpolation="None",
+            )
         ax2.axes.get_xaxis().set_visible(False)
         ax2.set_yticklabels([])
         ax2.set_yticks([])
         if i == 0:
-            ax2.set_ylabel('CENTER', fontdict=font)
+            ax2.set_ylabel("CENTER", fontdict=font)
         plt.tight_layout(pad=0, h_pad=0)
     plt.show()
 
 
-def visualize_crop_3d(data_dir, project_name, train_val_dir, center, new_cmp, anisotropy, index=None):
+def visualize_crop_3d(
+    data_dir, project_name, train_val_dir, center, new_cmp, anisotropy, index=None
+):
     """
-                Parameters
-                -------
-                data_dir: str
-                    Place where crops are (for example: data_dir='crops')
-                project_name: str
-                    Name of Project (for example: project_name= 'dsb-2018')
-                train_val_dir: str
-                    One of 'train' or 'val'
-                center: str
-                    One of 'medoid', 'centroid', 'approximate-medoid'
-                new_cmp: Color Map
-
-                anisotropy: float
-                    Ratio of voxel size along z dimension to size along x or y dimension
-                index: int
-                    Index of one of the images
-                Returns
-                -------
-
-                """
-
-    font = {'family': 'serif',
-            'color': 'black',
-            'weight': 'bold',
-            'size': 16,
-            }
-    im_filenames = sorted(glob(os.path.join(data_dir, project_name, train_val_dir) + '/images/*.tif'))
-    ma_filenames = sorted(glob(os.path.join(data_dir, project_name, train_val_dir) + '/masks/*.tif'))
+    Parameters
+    -------
+    data_dir: str
+        Place where crops are (for example: data_dir='crops')
+    project_name: str
+        Name of Project (for example: project_name= 'dsb-2018')
+    train_val_dir: str
+        One of 'train' or 'val'
+    center: str
+        One of 'medoid', 'centroid', 'approximate-medoid'
+    new_cmp: Color Map
+
+    anisotropy: float
+        Ratio of voxel size along z dimension to size along x or y dimension
+    index: int
+        Index of one of the images
+    Returns
+    -------
+
+    """
+
+    font = {
+        "family": "serif",
+        "color": "black",
+        "weight": "bold",
+        "size": 16,
+    }
+    im_filenames = sorted(
+        glob(os.path.join(data_dir, project_name, train_val_dir) + "/images/*.tif")
+    )
+    ma_filenames = sorted(
+        glob(os.path.join(data_dir, project_name, train_val_dir) + "/masks/*.tif")
+    )
     center_filenames = sorted(
-        glob(os.path.join(data_dir, project_name, train_val_dir) + '/center-' + center + '/*.tif'))
+        glob(
+            os.path.join(data_dir, project_name, train_val_dir)
+            + "/center-"
+            + center
+            + "/*.tif"
+        )
+    )
     if index is None:
         index = np.random.randint(0, len(im_filenames), 1)[0]
     fig = plt.figure(constrained_layout=False, figsize=(12, 10))
-    # rows represent three views of the sstack. cols represent the image, center and label
+    # rows represent three views of the sstack.
+    # cols represent the image, center and label
     spec = gridspec.GridSpec(ncols=3, nrows=3, figure=fig)
     im = tifffile.imread(im_filenames[index])
     label, _, _ = relabel_sequential(tifffile.imread(ma_filenames[index]))
@@ -256,67 +298,69 @@ def visualize_crop_3d(data_dir, project_name, train_val_dir, center, new_cmp, an
     x_mid = im.shape[2] // 2
 
     ax0 = fig.add_subplot(spec[0, 0])
-    ax0.imshow(im[z_mid, ...], cmap='magma', interpolation='None')
+    ax0.imshow(im[z_mid, ...], cmap="magma", interpolation="None")
     ax0.set_yticklabels([])
     ax0.set_yticks([])
     ax0.set_xticklabels([])
     ax0.set_xticks([])
-    ax0.set_ylabel('IM', fontdict=font)
-    ax0.set_xlabel('xy', fontdict=font)
-    ax0.xaxis.set_label_position('top')
+    ax0.set_ylabel("IM", fontdict=font)
+    ax0.set_xlabel("xy", fontdict=font)
+    ax0.xaxis.set_label_position("top")
 
     ax1 = fig.add_subplot(spec[1, 0])
-    ax1.imshow(label[z_mid, ...], cmap=new_cmp, interpolation='None')
+    ax1.imshow(label[z_mid, ...], cmap=new_cmp, interpolation="None")
     ax1.axes.get_xaxis().set_visible(False)
     ax1.set_yticklabels([])
     ax1.set_yticks([])
-    ax1.set_ylabel('MASK', fontdict=font)
+    ax1.set_ylabel("MASK", fontdict=font)
 
     ax2 = fig.add_subplot(spec[2, 0])
-    ax2.imshow(center_im[z_mid, ...], cmap='gray', interpolation='None')
+    ax2.imshow(center_im[z_mid, ...], cmap="gray", interpolation="None")
     ax2.axes.get_xaxis().set_visible(False)
     ax2.set_yticklabels([])
     ax2.set_yticks([])
-    ax2.set_ylabel('CENTER', fontdict=font)
+    ax2.set_ylabel("CENTER", fontdict=font)
 
     # use interpolation
     ax6 = fig.add_subplot(spec[0, 1])
-    ax6.imshow(im[..., x_mid], cmap='magma', interpolation='None')
+    ax6.imshow(im[..., x_mid], cmap="magma", interpolation="None")
     ax6.set_yticklabels([])
     ax6.set_yticks([])
     ax6.set_xticklabels([])
     ax6.set_xticks([])
-    ax6.set_xlabel('yz', fontdict=font)
-    ax6.xaxis.set_label_position('top')
+    ax6.set_xlabel("yz", fontdict=font)
+    ax6.xaxis.set_label_position("top")
 
     ax7 = fig.add_subplot(spec[1, 1])
-    ax7.imshow(label[..., x_mid], cmap=new_cmp, interpolation='None')
+    ax7.imshow(label[..., x_mid], cmap=new_cmp, interpolation="None")
     ax7.axes.get_xaxis().set_visible(False)
     ax7.set_yticklabels([])
     ax7.set_yticks([])
 
     ax8 = fig.add_subplot(spec[2, 1])
-    ax8.imshow(center_im[..., x_mid], cmap='gray', interpolation='None')
+    ax8.imshow(center_im[..., x_mid], cmap="gray", interpolation="None")
     ax8.axes.get_xaxis().set_visible(False)
     ax8.set_yticklabels([])
     ax8.set_yticks([])
 
     ax3 = fig.add_subplot(spec[0, 2])
-    ax3.imshow(np.transpose(im[:, y_mid, ...]), cmap='magma', interpolation='None')
+    ax3.imshow(np.transpose(im[:, y_mid, ...]), cmap="magma", interpolation="None")
     ax3.set_yticklabels([])
     ax3.set_yticks([])
     ax3.set_xticklabels([])
     ax3.set_xticks([])
-    ax3.set_xlabel('zx', fontdict=font)
-    ax3.xaxis.set_label_position('top')
+    ax3.set_xlabel("zx", fontdict=font)
+    ax3.xaxis.set_label_position("top")
 
     ax4 = fig.add_subplot(spec[1, 2])
-    ax4.imshow(np.transpose(label[:, y_mid, ...]), cmap=new_cmp, interpolation='None')
+    ax4.imshow(np.transpose(label[:, y_mid, ...]), cmap=new_cmp, interpolation="None")
     ax4.axes.get_xaxis().set_visible(False)
     ax4.set_yticklabels([])
     ax4.set_yticks([])
     ax5 = fig.add_subplot(spec[2, 2])
-    ax5.imshow(np.transpose(center_im[:, y_mid, ...]), cmap='gray', interpolation='None')
+    ax5.imshow(
+        np.transpose(center_im[:, y_mid, ...]), cmap="gray", interpolation="None"
+    )
     ax5.axes.get_xaxis().set_visible(False)
     ax5.set_yticklabels([])
     ax5.set_yticks([])
@@ -325,33 +369,37 @@ def visualize_crop_3d(data_dir, project_name, train_val_dir, center, new_cmp, an
     plt.show()
 
 
-def visualize_3d(im_filename, gt_filename, pred_filename, seed_filename, new_cmp, anisotropy):
+def visualize_3d(
+    im_filename, gt_filename, pred_filename, seed_filename, new_cmp, anisotropy
+):
     """
-        Parameters
-        -------
-        im_filename: str
+    Parameters
+    -------
+    im_filename: str
 
-        gt_filename: str
+    gt_filename: str
 
-        pred_filename: str
+    pred_filename: str
 
-        seed_filename: str
+    seed_filename: str
 
-        new_cmp: color map
+    new_cmp: color map
 
-        anisotropy: float
-            Ratio of voxel size along z dimension to size along x or y dimension
-        Returns
-        -------
+    anisotropy: float
+        Ratio of voxel size along z dimension to size along x or y dimension
+    Returns
+    -------
 
-        """
-    font = {'family': 'serif',
-            'color': 'black',
-            'weight': 'bold',
-            'size': 16,
-            }
+    """
+    font = {
+        "family": "serif",
+        "color": "black",
+        "weight": "bold",
+        "size": 16,
+    }
     fig = plt.figure(constrained_layout=False, figsize=(12, 10))
-    # rows represent three views of the sstack. cols represent the image, center and label
+    # rows represent three views of the sstack.
+    # cols represent the image, center and label
     spec = gridspec.GridSpec(ncols=3, nrows=4, figure=fig)
     im = tifffile.imread(im_filename)
     if gt_filename is not None:
@@ -370,90 +418,94 @@ def visualize_3d(im_filename, gt_filename, pred_filename, seed_filename, new_cmp
     x_mid = im.shape[2] // 2
 
     ax0 = fig.add_subplot(spec[0, 0])
-    ax0.imshow(im[z_mid, ...], cmap='magma', interpolation='None')
+    ax0.imshow(im[z_mid, ...], cmap="magma", interpolation="None")
     ax0.set_yticklabels([])
     ax0.set_yticks([])
     ax0.set_xticklabels([])
     ax0.set_xticks([])
-    ax0.set_ylabel('IM', fontdict=font)
-    ax0.set_xlabel('xy', fontdict=font)
-    ax0.xaxis.set_label_position('top')
+    ax0.set_ylabel("IM", fontdict=font)
+    ax0.set_xlabel("xy", fontdict=font)
+    ax0.xaxis.set_label_position("top")
 
     ax1 = fig.add_subplot(spec[1, 0])
     if gt_filename is not None:
-        ax1.imshow(gt_label[z_mid, ...], cmap=new_cmp, interpolation='None')
+        ax1.imshow(gt_label[z_mid, ...], cmap=new_cmp, interpolation="None")
     ax1.axes.get_xaxis().set_visible(False)
     ax1.set_yticklabels([])
     ax1.set_yticks([])
-    ax1.set_ylabel('GT', fontdict=font)
+    ax1.set_ylabel("GT", fontdict=font)
 
     ax2 = fig.add_subplot(spec[2, 0])
-    ax2.imshow(pred_label[z_mid, ...], cmap=new_cmp, interpolation='None')
+    ax2.imshow(pred_label[z_mid, ...], cmap=new_cmp, interpolation="None")
     ax2.axes.get_xaxis().set_visible(False)
     ax2.set_yticklabels([])
     ax2.set_yticks([])
-    ax2.set_ylabel('PRED', fontdict=font)
+    ax2.set_ylabel("PRED", fontdict=font)
 
     ax3 = fig.add_subplot(spec[3, 0])
-    ax3.imshow(seed[z_mid, ...], cmap='magma', interpolation='None')
+    ax3.imshow(seed[z_mid, ...], cmap="magma", interpolation="None")
     ax3.axes.get_xaxis().set_visible(False)
     ax3.set_yticklabels([])
     ax3.set_yticks([])
-    ax3.set_ylabel('SEED', fontdict=font)
+    ax3.set_ylabel("SEED", fontdict=font)
 
     # use interpolation
     ax8 = fig.add_subplot(spec[0, 1])
-    ax8.imshow(im[..., x_mid], cmap='magma', interpolation='None')
+    ax8.imshow(im[..., x_mid], cmap="magma", interpolation="None")
     ax8.set_yticklabels([])
     ax8.set_yticks([])
     ax8.set_xticklabels([])
     ax8.set_xticks([])
-    ax8.set_xlabel('yz', fontdict=font)
-    ax8.xaxis.set_label_position('top')
+    ax8.set_xlabel("yz", fontdict=font)
+    ax8.xaxis.set_label_position("top")
 
     ax9 = fig.add_subplot(spec[1, 1])
     if gt_filename is not None:
-        ax9.imshow(gt_label[..., x_mid], cmap=new_cmp, interpolation='None')
+        ax9.imshow(gt_label[..., x_mid], cmap=new_cmp, interpolation="None")
     ax9.axes.get_xaxis().set_visible(False)
     ax9.set_yticklabels([])
     ax9.set_yticks([])
 
     ax10 = fig.add_subplot(spec[2, 1])
-    ax10.imshow(pred_label[..., x_mid], cmap=new_cmp, interpolation='None')
+    ax10.imshow(pred_label[..., x_mid], cmap=new_cmp, interpolation="None")
     ax10.axes.get_xaxis().set_visible(False)
     ax10.set_yticklabels([])
     ax10.set_yticks([])
 
     ax11 = fig.add_subplot(spec[3, 1])
-    ax11.imshow(seed[..., x_mid], cmap='magma', interpolation='None')
+    ax11.imshow(seed[..., x_mid], cmap="magma", interpolation="None")
     ax11.axes.get_xaxis().set_visible(False)
     ax11.set_yticklabels([])
     ax11.set_yticks([])
 
     ax4 = fig.add_subplot(spec[0, 2])
-    ax4.imshow(np.transpose(im[:, y_mid, ...]), cmap='magma', interpolation='None')
+    ax4.imshow(np.transpose(im[:, y_mid, ...]), cmap="magma", interpolation="None")
     ax4.set_yticklabels([])
     ax4.set_yticks([])
     ax4.set_xticklabels([])
     ax4.set_xticks([])
-    ax4.set_xlabel('zx', fontdict=font)
-    ax4.xaxis.set_label_position('top')
+    ax4.set_xlabel("zx", fontdict=font)
+    ax4.xaxis.set_label_position("top")
 
     ax5 = fig.add_subplot(spec[1, 2])
     if gt_filename is not None:
-        ax5.imshow(np.transpose(gt_label[:, y_mid, ...]), cmap=new_cmp, interpolation='None')
+        ax5.imshow(
+            np.transpose(gt_label[:, y_mid, ...]), cmap=new_cmp, interpolation="None"
+        )
     ax5.axes.get_xaxis().set_visible(False)
     ax5.set_yticklabels([])
     ax5.set_yticks([])
 
     ax6 = fig.add_subplot(spec[2, 2])
-    ax6.imshow(np.transpose(pred_label[:, y_mid, ...]), cmap=new_cmp, interpolation='None')
+    ax6.imshow(
+        np.transpose(pred_label[:, y_mid, ...]), cmap=new_cmp, interpolation="None"
+    )
     ax6.axes.get_xaxis().set_visible(False)
     ax6.set_yticklabels([])
     ax6.set_yticks([])
 
     ax7 = fig.add_subplot(spec[3, 2])
-    ax7.imshow(np.transpose(seed[:, y_mid, ...]), cmap='magma', interpolation='None')
+    ax7.imshow(np.transpose(seed[:, y_mid, ...]), cmap="magma", interpolation="None")
     ax7.axes.get_xaxis().set_visible(False)
     ax7.set_yticklabels([])
     ax7.set_yticks([])
diff --git a/README.md b/README.md
index d116f27..e7c4ca5 100644
--- a/README.md
+++ b/README.md
@@ -24,18 +24,29 @@ With `EmbedSeg`, we obtain state-of-the-art results on multiple real-world micro
 
 
 ### Dependencies 
-We have tested this implementation using `pytorch` version 1.10.0 and `cudatoolkit` version 10.2 on a `linux` OS machine. 
-One could execute these lines of code to run this branch:
+
+One could execute these lines of code to run this branch with GPU support:
 
 ```
-conda create -n EmbedSegEnv python==3.7
-conda activate EmbedSegEnv
-conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
+mamba create -n EmbedSeg python
+mamba activate EmbedSeg
+mamba install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
 git clone https://github.com/juglab/EmbedSeg.git
 cd EmbedSeg
 pip install -e .
 ```
 
+For CPU support, one could execute the following lines of code:
+
+
+```
+mamba create -n EmbedSeg python
+mamba activate EmbedSeg
+pip install torch torchvision
+git clone https://github.com/juglab/EmbedSeg.git
+cd EmbedSeg
+```
+
 ### Getting Started
 
 Look in the `examples` directory,  and try out the `DSB-2018` notebooks for 2D images or `Mouse-Organoid-Cells-CBG` notebooks for volumetric (3D) images. Please make sure to select `Kernel > Change kernel` to `EmbedSegEnv`.   
diff --git a/setup.py b/setup.py
index 765c6c0..fd58b93 100644
--- a/setup.py
+++ b/setup.py
@@ -6,11 +6,12 @@
     long_description = fh.read()
 
 setup(
-    name="EmbedSeg", 
+    name="EmbedSeg",
     version="0.2.5",
     author="Manan Lalit, Pavel Tomancak, Florian Jug",
-    author_email="lalit@mpi-cbg.de",
-    description="EmbedSeg provides automatic detection and segmentation of objects in microscopy images",
+    author_email="lalitm@janelia.hhmi.org",
+    description="EmbedSeg provides automatic detection \
+            and segmentation of objects in microscopy images",
     long_description=long_description,
     long_description_content_type="text/markdown",
     url="https://github.com/juglab/EmbedSeg/",
@@ -18,11 +19,11 @@
         "Bug Tracker": "https://github.com/juglab/EmbedSeg/issues",
     },
     classifiers=[
-        'Programming Language :: Python',
-        'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3.7',
-        'Programming Language :: Python :: 3.8',
-        'Programming Language :: Python :: 3.9',
+        "Programming Language :: Python",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
         "License :: OSI Approved :: MIT License",
         "Operating System :: OS Independent",
     ],
@@ -30,20 +31,19 @@
     packages=setuptools.find_packages(),
     python_requires=">=3.7",
     install_requires=[
-          'matplotlib',
-          'numpy',
-          'scipy',
-          "tifffile",
-          "numba",
-          "tqdm",
-          "jupyter",
-          "pandas",
-          "seaborn",
-          "scikit-image",
-          "colorspacious",
-          "pytest",
-          "imagecodecs",
-          "pycocotools"
-        ]
+        "matplotlib",
+        "numpy",
+        "scipy",
+        "tifffile",
+        "numba",
+        "tqdm",
+        "jupyter",
+        "pandas",
+        "seaborn",
+        "scikit-image",
+        "colorspacious",
+        "pytest",
+        "imagecodecs",
+        "pycocotools",
+    ],
 )
-


Test patch:
diff --git a/EmbedSeg/_tests/test_center_crops.py b/EmbedSeg/_tests/test_center_crops.py
index e2bce04..af19095 100644
--- a/EmbedSeg/_tests/test_center_crops.py
+++ b/EmbedSeg/_tests/test_center_crops.py
@@ -1,23 +1,22 @@
-from EmbedSeg.utils.generate_crops import generate_center_image, generate_center_image_3d
-import pytest
+from EmbedSeg.utils.generate_crops import generate_center_image
 import numpy as np
 
+
 def test_generate_center_image_1():
     ma = np.zeros((100, 100))
-    ma [30:-30, 30:-30] = 2
+    ma[30:-30, 30:-30] = 2
     ids = np.unique(ma)
-    ids = ids[ids!=0]
-    center_image = generate_center_image(ma, center='centroid', ids=ids, one_hot = False )
-    y, x = np.where(center_image==True)
-    assert y==50 and x==50
+    ids = ids[ids != 0]
+    center_image = generate_center_image(ma, center="centroid", ids=ids, one_hot=False)
+    y, x = np.where(center_image is True)
+    assert y == 50 and x == 50
+
 
 def test_generate_center_image_2():
     ma = np.zeros((100, 100))
-    ma [30:-30, 30:-30] = 2
+    ma[30:-30, 30:-30] = 2
     ids = np.unique(ma)
-    ids = ids[ids!=0]
-    center_image = generate_center_image(ma, center='medoid', ids=ids, one_hot = False )
-    y, x = np.where(center_image==True)
-    assert np.abs(y-50)<=1 and np.abs(x-50)<=1    
-
-
+    ids = ids[ids != 0]
+    center_image = generate_center_image(ma, center="medoid", ids=ids, one_hot=False)
+    y, x = np.where(center_image is True)
+    assert np.abs(y - 50) <= 1 and np.abs(x - 50) <= 1
diff --git a/EmbedSeg/test.py b/EmbedSeg/test.py
index ae22d8c..0404bf0 100644
--- a/EmbedSeg/test.py
+++ b/EmbedSeg/test.py
@@ -8,8 +8,6 @@
 from EmbedSeg.models import get_model
 from EmbedSeg.utils.metrics import matching_dataset, obtain_APdsb_one_hot
 from EmbedSeg.utils.utils import Cluster, Cluster_3d
-
-torch.backends.cudnn.benchmark = True
 import numpy as np
 from tifffile import imsave
 from EmbedSeg.utils.test_time_augmentation import apply_tta_2d, apply_tta_3d
@@ -17,152 +15,304 @@
 from scipy.optimize import minimize_scalar, linear_sum_assignment
 from skimage.segmentation import relabel_sequential
 
+torch.backends.cudnn.benchmark = True
+
 
-def begin_evaluating(test_configs, optimize=False, maxiter=10, verbose=False, mask_region=None):
+def begin_evaluating(
+    test_configs, optimize=False, maxiter=10, verbose=False, mask_region=None
+):
     """Entry function for inferring on test images
 
-        Parameters
-        ----------
-        test_configs : dictionary
-            Dictionary containing testing-specific parameters (for e.g. the `seed_thresh`  to use)
-        optimize : bool, optional
-            It is possible to determine the best performing `fg_thresh` by optimizing over different values on the validation sub-set
-            By default and in the absence of optimization (i.e. `optimize=False`), the fg_thresh  is set equal to 0.5
-        maxiter: int
-            Number of iterations of optimization.
-            Comes into play, only if `optimize=True`
-        verbose: bool, optional
-            If set equal to True, prints the AP_dsb for each image individually
-        mask_region: list of lists, optional
-            If a certain region of the image is not labelled in the GT label mask, that can be specified here.
-            This enables comparison of the model prediction only with the area which is labeled in the GT label mask
-        Returns
-        -------
-        result_dic: Dictionary
-            Keys include the employed `fg_thresh` and the corresponding `AP_dsb` at IoU threshold = 0.5
-        """
-    n_sigma = test_configs['n_sigma']
-    ap_val = test_configs['ap_val']
-    min_mask_sum = test_configs['min_mask_sum']
-    min_unclustered_sum = test_configs['min_unclustered_sum']
-    min_object_size = test_configs['min_object_size']
-    mean_object_size = test_configs['mean_object_size']
-    tta = test_configs['tta']
-    seed_thresh = test_configs['seed_thresh']
-    fg_thresh = test_configs['fg_thresh']
-    save_images = test_configs['save_images']
-    save_results = test_configs['save_results']
-    save_dir = test_configs['save_dir']
-    anisotropy_factor = test_configs['anisotropy_factor']
-    grid_x = test_configs['grid_x']
-    grid_y = test_configs['grid_y']
-    grid_z = test_configs['grid_z']
-    pixel_x = test_configs['pixel_x']
-    pixel_y = test_configs['pixel_y']
-    pixel_z = test_configs['pixel_z']
-    one_hot = test_configs['dataset']['kwargs']['one_hot']
-    cluster_fast = test_configs['cluster_fast']
-    expand_grid = test_configs['expand_grid']
-    uniform_ds_factor = test_configs['uniform_ds_factor']
+    Parameters
+    ----------
+    test_configs : dictionary
+        Dictionary containing testing-specific parameters
+        (for e.g. the `seed_thresh`  to use)
+    optimize : bool, optional
+        It is possible to determine the best performing `fg_thresh`
+        by optimizing over different values on the validation sub-set
+        By default and in the absence of optimization (i.e. `optimize=False`),
+        the fg_thresh  is set equal to 0.5
+    maxiter: int
+        Number of iterations of optimization.
+        Comes into play, only if `optimize=True`
+    verbose: bool, optional
+        If set equal to True, prints the AP_dsb for each image individually
+    mask_region: list of lists, optional
+        If a certain region of the image != labelled in the GT label mask,
+        that can be specified here.
+        This enables comparison of the model prediction only
+        with the area which is labeled in the GT label mask
+    Returns
+    -------
+    result_dic: Dictionary
+        Keys include the employed `fg_thresh` and the corresponding `AP_dsb`
+        at IoU threshold = 0.5
+    """
+    n_sigma = test_configs["n_sigma"]
+    ap_val = test_configs["ap_val"]
+    min_mask_sum = test_configs["min_mask_sum"]
+    min_unclustered_sum = test_configs["min_unclustered_sum"]
+    min_object_size = test_configs["min_object_size"]
+    mean_object_size = test_configs["mean_object_size"]
+    tta = test_configs["tta"]
+    seed_thresh = test_configs["seed_thresh"]
+    fg_thresh = test_configs["fg_thresh"]
+    save_images = test_configs["save_images"]
+    save_results = test_configs["save_results"]
+    save_dir = test_configs["save_dir"]
+    anisotropy_factor = test_configs["anisotropy_factor"]
+    grid_x = test_configs["grid_x"]
+    grid_y = test_configs["grid_y"]
+    grid_z = test_configs["grid_z"]
+    pixel_x = test_configs["pixel_x"]
+    pixel_y = test_configs["pixel_y"]
+    pixel_z = test_configs["pixel_z"]
+    one_hot = test_configs["dataset"]["kwargs"]["one_hot"]
+    cluster_fast = test_configs["cluster_fast"]
+    expand_grid = test_configs["expand_grid"]
+    uniform_ds_factor = test_configs["uniform_ds_factor"]
 
     # set device
-    device = torch.device("cuda:0" if test_configs['cuda'] else "cpu")
+    device = torch.device(test_configs["device"])
 
     # dataloader
-    dataset = get_dataset(test_configs['dataset']['name'], test_configs['dataset']['kwargs'])
-    dataset_it = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, drop_last=False, num_workers=test_configs['num_workers'],
-                                             pin_memory=True if test_configs['cuda'] else False)
+    dataset = get_dataset(
+        test_configs["dataset"]["name"], test_configs["dataset"]["kwargs"]
+    )
+    dataset_it = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=1,
+        shuffle=False,
+        drop_last=False,
+        num_workers=test_configs["num_workers"],
+        pin_memory=True if test_configs["device"][:4] == "cuda" else False,
+    )
 
     # load model
-    model = get_model(test_configs['model']['name'], test_configs['model']['kwargs'])
-    model = torch.nn.DataParallel(model).to(device)
+    model = get_model(test_configs["model"]["name"], test_configs["model"]["kwargs"])
+    # model = torch.nn.DataParallel(model).to(device)
+    model = model.to(device)
 
     # load snapshot
-    if os.path.exists(test_configs['checkpoint_path']):
-        state = torch.load(test_configs['checkpoint_path'])
-        model.load_state_dict(state['model_state_dict'], strict=True)
+    if os.path.exists(test_configs["checkpoint_path"]):
+        state = torch.load(test_configs["checkpoint_path"], map_location=device)
+        model.load_state_dict(state["model_state_dict"], strict=True)  # TODO
     else:
-        assert False, 'checkpoint_path {} does not exist!'.format(test_configs['checkpoint_path'])
+        assert False, "checkpoint_path {} does not exist!".format(
+            test_configs["checkpoint_path"]
+        )
 
     # test on evaluation images:
     result_dic = {}
-    if (test_configs['name'] == '2d'):
-        args = (seed_thresh, ap_val, min_mask_sum, min_unclustered_sum, min_object_size, tta,
-                model, dataset_it, save_images, save_results, save_dir, verbose, grid_x,
-                grid_y, pixel_x, pixel_y, one_hot, n_sigma, cluster_fast, expand_grid)
+    if test_configs["name"] == "2d":
+        args = (
+            seed_thresh,
+            ap_val,
+            min_mask_sum,
+            min_unclustered_sum,
+            min_object_size,
+            tta,
+            model,
+            dataset_it,
+            save_images,
+            save_results,
+            save_dir,
+            verbose,
+            grid_x,
+            grid_y,
+            pixel_x,
+            pixel_y,
+            one_hot,
+            n_sigma,
+            cluster_fast,
+            expand_grid,
+            device,
+        )
         if optimize:
-            result = minimize_scalar(fun=test, bounds=(0.3, 0.7), method='bounded',
-                                     args=args, options={'maxiter': maxiter})
-            result_dic['seed_thresh'] = seed_thresh
-            result_dic['fg_thresh'] = result.x
-            result_dic['AP_dsb_05'] = -result.fun
-            print("Optimal fg thresh parameter calculated equal to {:.05f}".format(result.x))
-            print("AP_dsb at IoU threshold = {} with fg_thresh = {:.05f} equals {:.05f}".format(0.50, result.x,
-                                                                                                -result.fun))
+            result = minimize_scalar(
+                fun=test,
+                bounds=(0.3, 0.7),
+                method="bounded",
+                args=args,
+                options={"maxiter": maxiter},
+            )
+            result_dic["seed_thresh"] = seed_thresh
+            result_dic["fg_thresh"] = result.x
+            result_dic["AP_dsb_05"] = -result.fun
+            print(
+                "Optimal fg thresh parameter calculated equal to {:.05f}".format(
+                    result.x
+                )
+            )
+            print(
+                "AP_dsb at IoU threshold = {} with \
+                        fg_thresh = {:.05f} equals {:.05f}".format(
+                    0.50, result.x, -result.fun
+                )
+            )
         else:
             result = test(fg_thresh, *args)
-            result_dic['fg_thresh'] = fg_thresh
-            result_dic['AP_dsb_05'] = -result
-    elif (test_configs['name'] == '3d'):
+            result_dic["fg_thresh"] = fg_thresh
+            result_dic["AP_dsb_05"] = -result
+    elif test_configs["name"] == "3d":
         args = (
-            seed_thresh, ap_val, min_mask_sum, min_unclustered_sum, min_object_size, tta, model, dataset_it,
+            seed_thresh,
+            ap_val,
+            min_mask_sum,
+            min_unclustered_sum,
+            min_object_size,
+            tta,
+            model,
+            dataset_it,
             save_images,
-            save_results, save_dir, verbose, grid_x, grid_y, grid_z, pixel_x, pixel_y, pixel_z, one_hot, mask_region,
-            n_sigma, cluster_fast,
-            expand_grid, uniform_ds_factor)
+            save_results,
+            save_dir,
+            verbose,
+            grid_x,
+            grid_y,
+            grid_z,
+            pixel_x,
+            pixel_y,
+            pixel_z,
+            one_hot,
+            mask_region,
+            n_sigma,
+            cluster_fast,
+            expand_grid,
+            uniform_ds_factor,
+            device,
+        )
 
         if optimize:
-            result = minimize_scalar(fun=test_3d, bounds=(0.3, 0.7), method='bounded', args=args,
-                                     options={'maxiter': maxiter})
-            result_dic['seed_thresh'] = seed_thresh
-            result_dic['fg_thresh'] = result.x
-            result_dic['AP_dsb_05'] = -result.fun
-            print("Optimal fg thresh parameter calculated equal to {:.05f}".format(result.x))
-            print("AP_dsb at IoU threshold = {} with fg_thresh = {:.05f} equals {:.05f}".format(0.50, result.x,
-                                                                                                -result.fun))
+            result = minimize_scalar(
+                fun=test_3d,
+                bounds=(0.3, 0.7),
+                method="bounded",
+                args=args,
+                options={"maxiter": maxiter},
+            )
+            result_dic["seed_thresh"] = seed_thresh
+            result_dic["fg_thresh"] = result.x
+            result_dic["AP_dsb_05"] = -result.fun
+            print(
+                "Optimal fg thresh parameter calculated equal to {:.05f}".format(
+                    result.x
+                )
+            )
+            print(
+                "AP_dsb at IoU threshold = {} with \
+                        fg_thresh = {:.05f} equals {:.05f}".format(
+                    0.50, result.x, -result.fun
+                )
+            )
         else:
             result = test_3d(fg_thresh, *args)
-            result_dic['fg_thresh'] = fg_thresh
-            result_dic['AP_dsb_05'] = -result
-    elif (test_configs['name'] == '3d_sliced'):
-        args = (seed_thresh, ap_val, min_mask_sum, min_unclustered_sum, min_object_size, tta, model, dataset_it,
-                save_images, save_results, save_dir, verbose, grid_x, grid_y, grid_z,
-                pixel_x, pixel_y, pixel_z, one_hot, mask_region, n_sigma, anisotropy_factor)
+            result_dic["fg_thresh"] = fg_thresh
+            result_dic["AP_dsb_05"] = -result
+    elif test_configs["name"] == "3d_sliced":
+        args = (
+            seed_thresh,
+            ap_val,
+            min_mask_sum,
+            min_unclustered_sum,
+            min_object_size,
+            tta,
+            model,
+            dataset_it,
+            save_images,
+            save_results,
+            save_dir,
+            verbose,
+            grid_x,
+            grid_y,
+            grid_z,
+            pixel_x,
+            pixel_y,
+            pixel_z,
+            one_hot,
+            mask_region,
+            n_sigma,
+            anisotropy_factor,
+            device,
+        )
         if optimize:
-            result = minimize_scalar(fun=test_3d_sliced, bounds=(0.3, 0.7), method='bounded', args=args,
-                                     options={'maxiter': maxiter})
-            result_dic['seed_thresh'] = seed_thresh
-            result_dic['fg_thresh'] = result.x
-            result_dic['AP_dsb_05'] = -result.fun
+            result = minimize_scalar(
+                fun=test_3d_sliced,
+                bounds=(0.3, 0.7),
+                method="bounded",
+                args=args,
+                options={"maxiter": maxiter},
+            )
+            result_dic["seed_thresh"] = seed_thresh
+            result_dic["fg_thresh"] = result.x
+            result_dic["AP_dsb_05"] = -result.fun
         else:
             result = test_3d_sliced(fg_thresh, *args)
-            result_dic['fg_thresh'] = fg_thresh
-            result_dic['AP_dsb_05'] = -result
-    elif (test_configs['name'] == '3d_ilp'):
+            result_dic["fg_thresh"] = fg_thresh
+            result_dic["AP_dsb_05"] = -result
+    elif test_configs["name"] == "3d_ilp":
         args = (
-            seed_thresh, ap_val, min_mask_sum, min_unclustered_sum, min_object_size, mean_object_size, tta, model,
+            seed_thresh,
+            ap_val,
+            min_mask_sum,
+            min_unclustered_sum,
+            min_object_size,
+            mean_object_size,
+            tta,
+            model,
             dataset_it,
-            save_images, save_results, save_dir, verbose, grid_x, grid_y, grid_z,
-            pixel_x, pixel_y, pixel_z, one_hot, mask_region, n_sigma, anisotropy_factor)
+            save_images,
+            save_results,
+            save_dir,
+            verbose,
+            grid_x,
+            grid_y,
+            grid_z,
+            pixel_x,
+            pixel_y,
+            pixel_z,
+            one_hot,
+            mask_region,
+            n_sigma,
+            anisotropy_factor,
+            device,
+        )
         if optimize:
-            result = minimize_scalar(fun=test_3d_ilp, bounds=(0.3, 0.7), method='bounded', args=args,
-                                     options={'maxiter': maxiter})
-            result_dic['seed_thresh'] = seed_thresh
-            result_dic['fg_thresh'] = result.x
-            result_dic['AP_dsb_05'] = -result.fun
+            result = minimize_scalar(
+                fun=test_3d_ilp,
+                bounds=(0.3, 0.7),
+                method="bounded",
+                args=args,
+                options={"maxiter": maxiter},
+            )
+            result_dic["seed_thresh"] = seed_thresh
+            result_dic["fg_thresh"] = result.x
+            result_dic["AP_dsb_05"] = -result.fun
         else:
             result = test_3d_ilp(fg_thresh, *args)
-            result_dic['fg_thresh'] = fg_thresh
-            result_dic['AP_dsb_05'] = -result
+            result_dic["fg_thresh"] = fg_thresh
+            result_dic["AP_dsb_05"] = -result
     return result_dic
 
 
-def stitch_2d(instance_map_tile, instance_map_current, y_current=None, x_current=None, last=1, num_overlap_pixels=4):
+def stitch_2d(
+    instance_map_tile,
+    instance_map_current,
+    y_current=None,
+    x_current=None,
+    last=1,
+    num_overlap_pixels=4,
+    device=None,
+):
     """
-    Stitching instance segmentations together in case the full 2D image doesn't fit in one go, on the GPU
+    Stitching instance segmentations together in case the full 2D image
+    doesn't fit in one go, on the GPU
     This function is executed only if `expand_grid` is set to False
-    The key idea is we identify the unique ids in the instance_map_current and the tile, but only in the overlap region.
-    Then we look at the IoU of these. If there is more than 50 % IoU, then these are considered to be the same
+    The key idea is we identify the unique ids in the instance_map_current
+    and the tile, but only in the overlap region.
+    Then we look at the IoU of these. If there is more than 50 % IoU,
+    then these are considered to be the same
     Else, a new id is generated!
 
     Parameters
@@ -193,6 +343,9 @@ def stitch_2d(instance_map_tile, instance_map_current, y_current=None, x_current
         updated instance segmentation over the full image)
 
     """
+    if device is None:
+        device = torch.device("cpu")
+
     mask = instance_map_tile > 0  # foreground region, which has been clustered
 
     h_tile = instance_map_tile.shape[0]
@@ -210,17 +363,23 @@ def stitch_2d(instance_map_tile, instance_map_current, y_current=None, x_current
         instance_map_tile = instance_map_tile_sequential
 
         # next pad the tile so that it is aligned wrt the complete image
-        # note that doing the padding ensures that the instance_map_tile is the same size as the instance_map_current
+        # note that doing the padding ensures that the instance_map_tile
+        # is the same size as the instance_map_current
 
-        instance_map_tile = np.pad(instance_map_tile,
-                                   ((y_current, np.maximum(0, h_current - y_current - h_tile)),
-                                    (x_current, np.maximum(0, w_current - x_current - w_tile))))
+        instance_map_tile = np.pad(
+            instance_map_tile,
+            (
+                (y_current, np.maximum(0, h_current - y_current - h_tile)),
+                (x_current, np.maximum(0, w_current - x_current - w_tile)),
+            ),
+        )
 
         # ensure that it has the same shape as instance_map_current
         instance_map_tile = instance_map_tile[:h_current, :w_current]
 
-        mask_overlap = np.zeros_like(
-            instance_map_tile)  # this just identifies the region where the tile overlaps with the `instance_map_current`
+        mask_overlap = np.zeros_like(instance_map_tile)
+        # this just identifies the region where the tile
+        # overlaps with the `instance_map_current`
 
         if y_current == 0 and x_current == 0:
             ids_tile = np.unique(instance_map_tile)
@@ -229,12 +388,24 @@ def stitch_2d(instance_map_tile, instance_map_current, y_current=None, x_current
             last = len(ids_tile) + 1
         else:
             if x_current != 0 and y_current == 0:
-                mask_overlap[y_current:y_current + h_tile, x_current:x_current + num_overlap_pixels] = 1
+                mask_overlap[
+                    y_current : y_current + h_tile,
+                    x_current : x_current + num_overlap_pixels,
+                ] = 1
             elif x_current == 0 and y_current != 0:
-                mask_overlap[y_current:y_current + num_overlap_pixels, x_current:x_current + w_tile] = 1
+                mask_overlap[
+                    y_current : y_current + num_overlap_pixels,
+                    x_current : x_current + w_tile,
+                ] = 1
             elif x_current != 0 and y_current != 0:
-                mask_overlap[y_current:y_current + h_tile, x_current:x_current + num_overlap_pixels] = 1
-                mask_overlap[y_current:y_current + num_overlap_pixels, x_current:x_current + w_tile] = 1
+                mask_overlap[
+                    y_current : y_current + h_tile,
+                    x_current : x_current + num_overlap_pixels,
+                ] = 1
+                mask_overlap[
+                    y_current : y_current + num_overlap_pixels,
+                    x_current : x_current + w_tile,
+                ] = 1
 
             # identify ids in the complete tile, not just the overlap region,
             ids_tile_all = np.unique(instance_map_tile)
@@ -248,92 +419,132 @@ def stitch_2d(instance_map_tile, instance_map_current, y_current=None, x_current
             ids_tile_notin_overlap = np.setdiff1d(ids_tile_all, ids_tile_overlap)
 
             # identify ids in `instance_map_current` but only in the overlap region
-            instance_map_current_masked = torch.from_numpy(instance_map_current * mask_overlap).cuda()
+            instance_map_current_masked = torch.from_numpy(
+                instance_map_current * mask_overlap
+            ).to(device)
 
-            ids_current_overlap = torch.unique(instance_map_current_masked).cpu().detach().numpy()
+            ids_current_overlap = (
+                torch.unique(instance_map_current_masked).cpu().detach().numpy()
+            )
             ids_current_overlap = ids_current_overlap[ids_current_overlap != 0]
 
             IoU_table = np.zeros((len(ids_tile_overlap), len(ids_current_overlap)))
-            instance_map_tile_masked = torch.from_numpy(instance_map_tile * mask_overlap).cuda()
+            instance_map_tile_masked = torch.from_numpy(
+                instance_map_tile * mask_overlap
+            ).to(device)
 
             # rows are ids in tile, cols are ids in GT instance map
 
             for i, id_tile_overlap in enumerate(ids_tile_overlap):
                 for j, id_current_overlap in enumerate(ids_current_overlap):
-
-                    intersection = ((instance_map_tile_masked == id_tile_overlap)
-                                    & (instance_map_current_masked == id_current_overlap)).sum()
-                    union = ((instance_map_tile_masked == id_tile_overlap)
-                             | (instance_map_current_masked == id_current_overlap)).sum()
+                    intersection = (
+                        (instance_map_tile_masked == id_tile_overlap)
+                        & (instance_map_current_masked == id_current_overlap)
+                    ).sum()
+                    union = (
+                        (instance_map_tile_masked == id_tile_overlap)
+                        | (instance_map_current_masked == id_current_overlap)
+                    ).sum()
                     if union != 0:
                         IoU_table[i, j] = intersection / union
                     else:
                         IoU_table[i, j] = 0.0
 
             row_indices, col_indices = linear_sum_assignment(-IoU_table)
-            matched_indices = np.array(list(zip(row_indices, col_indices)))  # list of (row, col) tuples
-            unmatched_indices_tile = np.setdiff1d(np.arange(len(ids_tile_overlap)), row_indices)
+            matched_indices = np.array(
+                list(zip(row_indices, col_indices))
+            )  # list of (row, col) tuples
+            unmatched_indices_tile = np.setdiff1d(
+                np.arange(len(ids_tile_overlap)), row_indices
+            )
 
             for m in matched_indices:
-                if (IoU_table[m[0], m[1]] >= 0.5):  # (tile, current)
-                    # wherever the tile is m[0], it should be assigned m[1] in the larger prediction image
-                    instance_map_current[instance_map_tile == ids_tile_overlap[m[0]]] = ids_current_overlap[m[1]]
-                elif (IoU_table[m[0], m[1]] == 0):
+                if IoU_table[m[0], m[1]] >= 0.5:  # (tile, current)
+                    # wherever the tile is m[0], it should be assigned
+                    # m[1] in the larger prediction image
+                    instance_map_current[
+                        instance_map_tile == ids_tile_overlap[m[0]]
+                    ] = ids_current_overlap[m[1]]
+                elif IoU_table[m[0], m[1]] == 0:
                     # there is no intersection
-                    instance_map_current[instance_map_tile == ids_tile_overlap[m[0]]] = last
+                    instance_map_current[
+                        instance_map_tile == ids_tile_overlap[m[0]]
+                    ] = last
                     last += 1
                 else:
                     # otherwise just take a union of the both ...
-                    # basically this should spawn a new label, since there was not a satisfactory match with any pre-existing id in the instance_map_current
-                    instance_map_current[instance_map_tile == ids_tile_overlap[m[0]]] = last
-                    # instance_map_current[instance_map_current == ids_current_overlap[m[1]]] = last  # TODO
+                    # basically this should spawn a new label, since there
+                    # was not a satisfactory match with any pre-existing id
+                    # in the instance_map_current
+                    instance_map_current[
+                        instance_map_tile == ids_tile_overlap[m[0]]
+                    ] = last
+                    # instance_map_current[instance_map_current == \
+                    #        ids_current_overlap[m[1]]] = last  # TODO
                     last += 1
             for index in unmatched_indices_tile:  # not a tuple
-                instance_map_current[instance_map_tile == ids_tile_overlap[index]] = last
+                instance_map_current[
+                    instance_map_tile == ids_tile_overlap[index]
+                ] = last
                 last += 1
             for id in ids_tile_notin_overlap:
                 instance_map_current[instance_map_tile == id] = last
                 last += 1
         return last, instance_map_current
     else:
-        return last, instance_map_current  # if there are no ids in tile, then just return
-
-
-def stitch_3d(instance_map_tile, instance_map_current, z_tile=None, y_tile=None, x_tile=None, last=1,
-              num_overlap_pixels=4):
-    """Stitching instance segmentations together in case the full 3D image doesn't fit in one go, on the GPU
-       This function is executed only if `expand_grid` is set to False
+        return (
+            last,
+            instance_map_current,
+        )  # if there are no ids in tile, then just return
+
+
+def stitch_3d(
+    instance_map_tile,
+    instance_map_current,
+    z_tile=None,
+    y_tile=None,
+    x_tile=None,
+    last=1,
+    num_overlap_pixels=4,
+    device=None,
+):
+    """Stitching instance segmentations together in case the
+    full 3D image doesn't fit in one go, on the GPU
+    This function is executed only if `expand_grid` is set to False
 
-                Parameters
-                ----------
-                instance_map_tile : numpy array
-                    instance segmentation over a tiled view of the image
+             Parameters
+             ----------
+             instance_map_tile : numpy array
+                 instance segmentation over a tiled view of the image
 
-                instance_map_current: numpy array
-                    instance segmentation over the complete, large image
+             instance_map_current: numpy array
+                 instance segmentation over the complete, large image
 
-                z_tile: int
-                    z position of the top left corner of the tile wrt the complete image
+             z_tile: int
+                 z position of the top left corner of the tile wrt the complete image
 
-                y_tile: int
-                    y position of the top left corner of the tile wrt the complete image
+             y_tile: int
+                 y position of the top left corner of the tile wrt the complete image
 
-                x_tile: int
-                    x position of the top left corner of the tile wrt the complete image
+             x_tile: int
+                 x position of the top left corner of the tile wrt the complete image
 
-                last: int
-                    number of objects currently present in the `instance_map_current`
+             last: int
+                 number of objects currently present in the `instance_map_current`
 
-                num_overlap_pixels: int
-                    number of overlapping pixels while considering the next tile
+             num_overlap_pixels: int
+                 number of overlapping pixels while considering the next tile
 
-                Returns
-                -------
-                tuple (int, numpy array)
-                    (updated number of objects currently present in the `instance_map_current`,
-                    updated instance segmentation over the full image)
+             Returns
+             -------
+             tuple (int, numpy array)
+                 (updated number of objects currently
+                  present in the `instance_map_current`,
+                 updated instance segmentation over the full image)
 
-                """
+    """
+    if device is None:
+        device = torch.device("cpu")
 
     mask = instance_map_tile > 0
 
@@ -352,42 +563,100 @@ def stitch_3d(instance_map_tile, instance_map_current, z_tile=None, y_tile=None,
 
         # next pad the tile so that it is aligned wrt the complete image
 
-        instance_map_tile = np.pad(instance_map_tile,
-                                   ((z_tile, np.maximum(0, instance_map_current.shape[0] - z_tile - D)),
-                                    (y_tile, np.maximum(0, instance_map_current.shape[1] - y_tile - H)),
-                                    (x_tile, np.maximum(0, instance_map_current.shape[2] - x_tile - W))))
+        instance_map_tile = np.pad(
+            instance_map_tile,
+            (
+                (z_tile, np.maximum(0, instance_map_current.shape[0] - z_tile - D)),
+                (y_tile, np.maximum(0, instance_map_current.shape[1] - y_tile - H)),
+                (x_tile, np.maximum(0, instance_map_current.shape[2] - x_tile - W)),
+            ),
+        )
 
         # ensure that it has the same shape as instance_map_current
-        instance_map_tile = instance_map_tile[:instance_map_current.shape[0], :instance_map_current.shape[1],
-                            :instance_map_current.shape[2]]
+        instance_map_tile = instance_map_tile[
+            : instance_map_current.shape[0],
+            : instance_map_current.shape[1],
+            : instance_map_current.shape[2],
+        ]
         mask_overlap = np.zeros_like(instance_map_tile)
 
         if z_tile == 0 and y_tile == 0 and x_tile == 0:
             ids_tile = np.unique(instance_map_tile)
             ids_tile = ids_tile[ids_tile != 0]
-            instance_map_current[:instance_map_tile.shape[0], :instance_map_tile.shape[1],
-            :instance_map_tile.shape[2]] = instance_map_tile
+            instance_map_current[
+                : instance_map_tile.shape[0],
+                : instance_map_tile.shape[1],
+                : instance_map_tile.shape[2],
+            ] = instance_map_tile
             last = len(ids_tile) + 1
         else:
             if x_tile != 0 and y_tile == 0 and z_tile == 0:
-                mask_overlap[z_tile:z_tile + D, y_tile:y_tile + H, x_tile:x_tile + num_overlap_pixels] = 1
+                mask_overlap[
+                    z_tile : z_tile + D,
+                    y_tile : y_tile + H,
+                    x_tile : x_tile + num_overlap_pixels,
+                ] = 1
             elif x_tile == 0 and y_tile != 0 and z_tile == 0:
-                mask_overlap[z_tile:z_tile + D, y_tile:y_tile + num_overlap_pixels, x_tile:x_tile + W] = 1
+                mask_overlap[
+                    z_tile : z_tile + D,
+                    y_tile : y_tile + num_overlap_pixels,
+                    x_tile : x_tile + W,
+                ] = 1
             elif x_tile != 0 and y_tile != 0 and z_tile == 0:
-                mask_overlap[z_tile:z_tile + D, y_tile:y_tile + H, x_tile:x_tile + num_overlap_pixels] = 1
-                mask_overlap[z_tile:z_tile + D, y_tile:y_tile + num_overlap_pixels, x_tile:x_tile + W] = 1
+                mask_overlap[
+                    z_tile : z_tile + D,
+                    y_tile : y_tile + H,
+                    x_tile : x_tile + num_overlap_pixels,
+                ] = 1
+                mask_overlap[
+                    z_tile : z_tile + D,
+                    y_tile : y_tile + num_overlap_pixels,
+                    x_tile : x_tile + W,
+                ] = 1
             elif x_tile == 0 and y_tile == 0 and z_tile != 0:
-                mask_overlap[z_tile:z_tile + num_overlap_pixels, y_tile:y_tile + H, x_tile:x_tile + W] = 1
+                mask_overlap[
+                    z_tile : z_tile + num_overlap_pixels,
+                    y_tile : y_tile + H,
+                    x_tile : x_tile + W,
+                ] = 1
             elif x_tile != 0 and y_tile == 0 and z_tile != 0:
-                mask_overlap[z_tile:z_tile + num_overlap_pixels, y_tile:y_tile + H, x_tile:x_tile + W] = 1
-                mask_overlap[z_tile:z_tile + D, y_tile:y_tile + H, x_tile:x_tile + num_overlap_pixels] = 1
+                mask_overlap[
+                    z_tile : z_tile + num_overlap_pixels,
+                    y_tile : y_tile + H,
+                    x_tile : x_tile + W,
+                ] = 1
+                mask_overlap[
+                    z_tile : z_tile + D,
+                    y_tile : y_tile + H,
+                    x_tile : x_tile + num_overlap_pixels,
+                ] = 1
             elif x_tile == 0 and y_tile != 0 and z_tile != 0:
-                mask_overlap[z_tile:z_tile + num_overlap_pixels, y_tile:y_tile + H, x_tile:x_tile + W] = 1
-                mask_overlap[z_tile:z_tile + D, y_tile:y_tile + num_overlap_pixels, x_tile:x_tile + W] = 1
+                mask_overlap[
+                    z_tile : z_tile + num_overlap_pixels,
+                    y_tile : y_tile + H,
+                    x_tile : x_tile + W,
+                ] = 1
+                mask_overlap[
+                    z_tile : z_tile + D,
+                    y_tile : y_tile + num_overlap_pixels,
+                    x_tile : x_tile + W,
+                ] = 1
             elif x_tile != 0 and y_tile != 0 and z_tile != 0:
-                mask_overlap[z_tile:z_tile + D, y_tile:y_tile + H, x_tile:x_tile + num_overlap_pixels] = 1
-                mask_overlap[z_tile:z_tile + D, y_tile:y_tile + num_overlap_pixels, x_tile:x_tile + W] = 1
-                mask_overlap[z_tile:z_tile + num_overlap_pixels, y_tile:y_tile + H, x_tile:x_tile + W] = 1
+                mask_overlap[
+                    z_tile : z_tile + D,
+                    y_tile : y_tile + H,
+                    x_tile : x_tile + num_overlap_pixels,
+                ] = 1
+                mask_overlap[
+                    z_tile : z_tile + D,
+                    y_tile : y_tile + num_overlap_pixels,
+                    x_tile : x_tile + W,
+                ] = 1
+                mask_overlap[
+                    z_tile : z_tile + num_overlap_pixels,
+                    y_tile : y_tile + H,
+                    x_tile : x_tile + W,
+                ] = 1
 
             # identify ids in the complete tile, not just the overlap region,
             ids_tile_all = np.unique(instance_map_tile)
@@ -401,47 +670,70 @@ def stitch_3d(instance_map_tile, instance_map_current, z_tile=None, y_tile=None,
             ids_tile_notin_overlap = np.setdiff1d(ids_tile_all, ids_tile_overlap)
 
             # identify ids in `instance_map_current` but only in the overlap region
-            instance_map_current_masked = torch.from_numpy(instance_map_current * mask_overlap).cuda()
+            instance_map_current_masked = torch.from_numpy(
+                instance_map_current * mask_overlap
+            ).to(device)
 
-            ids_current = torch.unique(instance_map_current_masked).cpu().detach().numpy()
+            ids_current = (
+                torch.unique(instance_map_current_masked).cpu().detach().numpy()
+            )
             ids_current = ids_current[ids_current != 0]
 
             IoU_table = np.zeros((len(ids_tile_overlap), len(ids_current)))
-            instance_map_tile_masked = torch.from_numpy(instance_map_tile * mask_overlap).cuda()
+            instance_map_tile_masked = torch.from_numpy(
+                instance_map_tile * mask_overlap
+            ).to(device)
 
             # rows are ids in tile, cols are ids in GT instance map
 
             for i, id_tile in enumerate(ids_tile_overlap):
                 for j, id_current in enumerate(ids_current):
-
-                    intersection = ((instance_map_tile_masked == id_tile)
-                                    & (instance_map_current_masked == id_current)).sum()
-                    union = ((instance_map_tile_masked == id_tile)
-                             | (instance_map_current_masked == id_current)).sum()
+                    intersection = (
+                        (instance_map_tile_masked == id_tile)
+                        & (instance_map_current_masked == id_current)
+                    ).sum()
+                    union = (
+                        (instance_map_tile_masked == id_tile)
+                        | (instance_map_current_masked == id_current)
+                    ).sum()
                     if union != 0:
                         IoU_table[i, j] = intersection / union
                     else:
                         IoU_table[i, j] = 0.0
 
             row_indices, col_indices = linear_sum_assignment(-IoU_table)
-            matched_indices = np.array(list(zip(row_indices, col_indices)))  # list of (row, col) tuples
-            unmatched_indices_tile = np.setdiff1d(np.arange(len(ids_tile_overlap)), row_indices)
+            matched_indices = np.array(
+                list(zip(row_indices, col_indices))
+            )  # list of (row, col) tuples
+            unmatched_indices_tile = np.setdiff1d(
+                np.arange(len(ids_tile_overlap)), row_indices
+            )
 
             for m in matched_indices:
-                if (IoU_table[m[0], m[1]] >= 0.5):  # (tile, current)
-                    # wherever the tile is m[0], it should be assigned m[1] in the larger prediction image
-                    instance_map_current[instance_map_tile == ids_tile_overlap[m[0]]] = ids_current[m[1]]
-                elif (IoU_table[m[0], m[1]] == 0):
+                if IoU_table[m[0], m[1]] >= 0.5:  # (tile, current)
+                    # wherever the tile is m[0],
+                    # it should be assigned m[1] in the larger prediction image
+                    instance_map_current[
+                        instance_map_tile == ids_tile_overlap[m[0]]
+                    ] = ids_current[m[1]]
+                elif IoU_table[m[0], m[1]] == 0:
                     # there is no intersection
-                    instance_map_current[instance_map_tile == ids_tile_overlap[m[0]]] = last
+                    instance_map_current[
+                        instance_map_tile == ids_tile_overlap[m[0]]
+                    ] = last
                     last += 1
                 else:
                     # otherwise just take a union of the both ...
-                    instance_map_current[instance_map_tile == ids_tile_overlap[m[0]]] = last
-                    # instance_map_current[instance_map_current == ids_current[m[1]]] = last
+                    instance_map_current[
+                        instance_map_tile == ids_tile_overlap[m[0]]
+                    ] = last
+                    # instance_map_current[instance_map_current
+                    #                       == ids_current[m[1]]] = last
                     last += 1
             for index in unmatched_indices_tile:  # not a tuple
-                instance_map_current[instance_map_tile == ids_tile_overlap[index]] = last
+                instance_map_current[
+                    instance_map_tile == ids_tile_overlap[index]
+                ] = last
                 last += 1
             for id in ids_tile_notin_overlap:
                 instance_map_current[instance_map_tile == id] = last
@@ -449,12 +741,26 @@ def stitch_3d(instance_map_tile, instance_map_current, z_tile=None, y_tile=None,
 
         return last, instance_map_current
     else:
-        return last, instance_map_current  # if there are no ids in tile, then just return
-
-
-def predict(im, model, tta, cluster_fast, n_sigma, fg_thresh, seed_thresh, min_mask_sum, min_unclustered_sum,
-            min_object_size,
-            cluster):
+        return (
+            last,
+            instance_map_current,
+        )  # if there are no ids in tile, then just return
+
+
+def predict(
+    im,
+    model,
+    tta,
+    cluster_fast,
+    n_sigma,
+    fg_thresh,
+    seed_thresh,
+    min_mask_sum,
+    min_unclustered_sum,
+    min_object_size,
+    cluster,
+    device,
+):
     """
 
     Parameters
@@ -476,7 +782,8 @@ def predict(im, model, tta, cluster_fast, n_sigma, fg_thresh, seed_thresh, min_m
     seed_thresh: float
         This should be set equal to `0.9` by default
     min_mask_sum: int
-        Only start creating instances, if there are at least `min_mask_sum` pixels in foreground!
+        Only start creating instances,
+        if there are at least `min_mask_sum` pixels in foreground!
     min_unclustered_sum: int
         Stop when the number of seed candidates are less than `min_unclustered_sum`
     min_object_size: int
@@ -503,43 +810,63 @@ def predict(im, model, tta, cluster_fast, n_sigma, fg_thresh, seed_thresh, min_m
         diff_x = 8 * (multiple_x + 1) - im.shape[3]
     else:
         diff_x = 0
-    p2d = (diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2)  # last dim, second last dim
+    p2d = (
+        diff_x // 2,
+        diff_x - diff_x // 2,
+        diff_y // 2,
+        diff_y - diff_y // 2,
+    )  # last dim, second last dim
 
     im = F.pad(im, p2d, "reflect")
 
-    if (tta):
+    if tta:
         output = apply_tta_2d(im, model)
     else:
         output = model(im)
     if cluster_fast:
-        instance_map = cluster.cluster(output[0],
-                                       n_sigma=n_sigma,
-                                       seed_thresh=seed_thresh,
-                                       min_mask_sum=min_mask_sum,
-                                       min_unclustered_sum=min_unclustered_sum,
-                                       min_object_size=min_object_size)
+        instance_map = cluster.cluster(
+            output[0],
+            n_sigma=n_sigma,
+            seed_thresh=seed_thresh,
+            min_mask_sum=min_mask_sum,
+            min_unclustered_sum=min_unclustered_sum,
+            min_object_size=min_object_size,
+        )
     else:
-        instance_map = cluster.cluster_local_maxima(output[0],
-                                                    n_sigma=n_sigma,
-                                                    fg_thresh=fg_thresh,
-                                                    min_mask_sum=min_mask_sum,
-                                                    min_unclustered_sum=min_unclustered_sum,
-                                                    min_object_size=min_object_size)
+        instance_map = cluster.cluster_local_maxima(
+            output[0],
+            n_sigma=n_sigma,
+            fg_thresh=fg_thresh,
+            min_mask_sum=min_mask_sum,
+            min_unclustered_sum=min_unclustered_sum,
+            min_object_size=min_object_size,
+        )
 
     seed_map = torch.sigmoid(output[0, -1, ...])
     # unpad instance_map, seed_map
-    if (diff_y - diff_y // 2) is not 0:
-        instance_map = instance_map[diff_y // 2:-(diff_y - diff_y // 2), ...]
-        seed_map = seed_map[diff_y // 2:-(diff_y - diff_y // 2), ...]
-    if (diff_x - diff_x // 2) is not 0:
-        instance_map = instance_map[..., diff_x // 2:-(diff_x - diff_x // 2)]
-        seed_map = seed_map[..., diff_x // 2:-(diff_x - diff_x // 2)]
+    if (diff_y - diff_y // 2) != 0:
+        instance_map = instance_map[diff_y // 2 : -(diff_y - diff_y // 2), ...]
+        seed_map = seed_map[diff_y // 2 : -(diff_y - diff_y // 2), ...]
+    if (diff_x - diff_x // 2) != 0:
+        instance_map = instance_map[..., diff_x // 2 : -(diff_x - diff_x // 2)]
+        seed_map = seed_map[..., diff_x // 2 : -(diff_x - diff_x // 2)]
     return instance_map, seed_map
 
 
-def predict_3d(im, model, tta, cluster_fast, n_sigma, fg_thresh, seed_thresh, min_mask_sum, min_unclustered_sum,
-            min_object_size,
-            cluster):
+def predict_3d(
+    im,
+    model,
+    tta,
+    cluster_fast,
+    n_sigma,
+    fg_thresh,
+    seed_thresh,
+    min_mask_sum,
+    min_unclustered_sum,
+    min_object_size,
+    cluster,
+    device,
+):
     """
 
     Parameters
@@ -561,7 +888,8 @@ def predict_3d(im, model, tta, cluster_fast, n_sigma, fg_thresh, seed_thresh, mi
     seed_thresh: float
         This should be set equal to `0.9` by default
     min_mask_sum: int
-        Only start creating instances, if there are at least `min_mask_sum` pixels in foreground!
+        Only start creating instances,
+        if there are at least `min_mask_sum` pixels in foreground!
     min_unclustered_sum: int
         Stop when the number of seed candidates are less than `min_unclustered_sum`
     min_object_size: int
@@ -583,26 +911,34 @@ def predict_3d(im, model, tta, cluster_fast, n_sigma, fg_thresh, seed_thresh, mi
             if iter == 0:
                 output_average = apply_tta_3d(im, model, iter)
             else:
-                output_average = 1 / (iter + 1) * (
-                        output_average * iter + apply_tta_3d(im, model, iter))  # iter
-        output = torch.from_numpy(output_average).float().cuda()
+                output_average = (
+                    1
+                    / (iter + 1)
+                    * (output_average * iter + apply_tta_3d(im, model, iter))
+                )  # iter
+        output = torch.from_numpy(output_average).float().to(device)
     else:
         output = model(im)
 
     if cluster_fast:
-        instance_map = cluster.cluster(output[0], n_sigma=n_sigma,
-                                       fg_thresh=fg_thresh,
-                                       seed_thresh=seed_thresh,
-                                       min_mask_sum=min_mask_sum,
-                                       min_unclustered_sum=min_unclustered_sum,
-                                       min_object_size=min_object_size,
-                                       )
+        instance_map = cluster.cluster(
+            output[0],
+            n_sigma=n_sigma,
+            fg_thresh=fg_thresh,
+            seed_thresh=seed_thresh,
+            min_mask_sum=min_mask_sum,
+            min_unclustered_sum=min_unclustered_sum,
+            min_object_size=min_object_size,
+        )
     else:
-        instance_map = cluster.cluster_local_maxima(output[0], n_sigma=n_sigma,
-                                                    fg_thresh=fg_thresh,
-                                                    min_mask_sum=min_mask_sum,
-                                                    min_unclustered_sum=min_unclustered_sum,
-                                                    min_object_size=min_object_size)
+        instance_map = cluster.cluster_local_maxima(
+            output[0],
+            n_sigma=n_sigma,
+            fg_thresh=fg_thresh,
+            min_mask_sum=min_mask_sum,
+            min_unclustered_sum=min_unclustered_sum,
+            min_object_size=min_object_size,
+        )
     seed_map = torch.sigmoid(output[0, -1, ...])
     # unpad instance_map, seed_map
 
@@ -618,141 +954,252 @@ def predict_3d(im, model, tta, cluster_fast, n_sigma, fg_thresh, seed_thresh, mi
     return instance_map, seed_map
 
 
-
 def test(fg_thresh, *args):
     """Infer the trained 2D model on 2D images
 
-            Parameters
-            ----------
-            fg_thresh : float
-                foreground threshold decides which pixels are considered for clustering, based on the predicted seediness scores at these pixels.
-            args: dictionary
-                Contains other paremeters such as `ap_val`, `seed_thresh` etc
-            Returns
-            -------
-            float
-                Average `AP_dsb` over all test images
-            """
-    seed_thresh, ap_val, min_mask_sum, min_unclustered_sum, min_object_size, tta, model, dataset_it, save_images, \
-    save_results, save_dir, verbose, grid_x, grid_y, pixel_x, pixel_y, one_hot, n_sigma, cluster_fast, expand_grid = args
+    Parameters
+    ----------
+    fg_thresh : float
+        foreground threshold decides which pixels are considered for clustering,
+        based on the predicted seediness scores at these pixels.
+    args: dictionary
+        Contains other paremeters such as `ap_val`, `seed_thresh` etc
+    Returns
+    -------
+    float
+        Average `AP_dsb` over all test images
+    """
+    (
+        seed_thresh,
+        ap_val,
+        min_mask_sum,
+        min_unclustered_sum,
+        min_object_size,
+        tta,
+        model,
+        dataset_it,
+        save_images,
+        save_results,
+        save_dir,
+        verbose,
+        grid_x,
+        grid_y,
+        pixel_x,
+        pixel_y,
+        one_hot,
+        n_sigma,
+        cluster_fast,
+        expand_grid,
+        device,
+    ) = args
 
     model.eval()
 
     # cluster module
-    cluster = Cluster(grid_y, grid_x, pixel_y, pixel_x)
+    cluster = Cluster(grid_y, grid_x, pixel_y, pixel_x, device)
 
     with torch.no_grad():
         result_list = []
         image_file_names = []
         for sample in tqdm(dataset_it):
+            im = sample["image"]  # B 1 Y X
+            im = im.to(device)
 
-            im = sample['image']  # B 1 Y X
             H, W = im.shape[2], im.shape[3]
             if H > grid_y or W > grid_x:
                 if expand_grid:
-                    # simple trick to expand the grid while keeping pixel resolution the same as before
+                    # simple trick to expand the grid while keeping
+                    # pixel resolution the same as before
                     H_, W_ = round_up_8(H), round_up_8(W)
                     temp = np.maximum(H_, W_)
                     H_ = temp
                     W_ = temp
                     pixel_x_modified = pixel_y_modified = H_ / grid_y
-                    cluster = Cluster(H_, W_, pixel_y_modified, pixel_x_modified)
-                    instance_map, seed_map = predict(im, model, tta, cluster_fast,
-                                                     n_sigma, fg_thresh, seed_thresh, min_mask_sum, min_unclustered_sum,
-                                                     min_object_size, cluster)
-
+                    cluster = Cluster(
+                        H_, W_, pixel_y_modified, pixel_x_modified, device
+                    )
+                    instance_map, seed_map = predict(
+                        im,
+                        model,
+                        tta,
+                        cluster_fast,
+                        n_sigma,
+                        fg_thresh,
+                        seed_thresh,
+                        min_mask_sum,
+                        min_unclustered_sum,
+                        min_object_size,
+                        cluster,
+                        device,
+                    )
 
                 else:
                     # here, we try stitching predictions instead
                     last = 1
                     instance_map = np.zeros((H, W), dtype=np.int16)
-                    seed_map = np.zeros((H, W), dtype=np.float)
+                    seed_map = np.zeros((H, W), dtype=float)
                     num_overlap_pixels = 4
                     for y in range(0, H, grid_y - num_overlap_pixels):
                         for x in range(0, W, grid_x - num_overlap_pixels):
-                            instance_map_tile, seed_map_tile = predict(im[:, :, y:y + grid_y, x:x + grid_x], model, tta,
-                                                                       cluster_fast,
-                                                                       n_sigma, fg_thresh, seed_thresh, min_mask_sum,
-                                                                       min_unclustered_sum, min_object_size,
-                                                                       cluster)
-                            last, instance_map = stitch_2d(instance_map_tile.cpu().detach().numpy(), instance_map, y, x,
-                                                           last, num_overlap_pixels)
-                            seed_map[y:y + grid_y, x:x + grid_x] = seed_map_tile.cpu().detach().numpy()
-                    instance_map = torch.from_numpy(instance_map).cuda()
-                    seed_map = torch.from_numpy(seed_map).float().cuda()
+                            instance_map_tile, seed_map_tile = predict(
+                                im[:, :, y : y + grid_y, x : x + grid_x],
+                                model,
+                                tta,
+                                cluster_fast,
+                                n_sigma,
+                                fg_thresh,
+                                seed_thresh,
+                                min_mask_sum,
+                                min_unclustered_sum,
+                                min_object_size,
+                                cluster,
+                                device,
+                            )
+                            last, instance_map = stitch_2d(
+                                instance_map_tile.cpu().detach().numpy(),
+                                instance_map,
+                                y,
+                                x,
+                                last,
+                                num_overlap_pixels,
+                                device,
+                            )
+                            seed_map[y : y + grid_y, x : x + grid_x] = (
+                                seed_map_tile.cpu().detach().numpy()
+                            )
+                    instance_map = torch.from_numpy(instance_map).to(device)
+                    seed_map = torch.from_numpy(seed_map).float().to(device)
             else:
-                instance_map, seed_map = predict(im, model, tta, cluster_fast,
-                                                 n_sigma, fg_thresh, seed_thresh, min_mask_sum, min_unclustered_sum,
-                                                 min_object_size, cluster)
-
-            base, _ = os.path.splitext(os.path.basename(sample['im_name'][0]))
+                instance_map, seed_map = predict(
+                    im,
+                    model,
+                    tta,
+                    cluster_fast,
+                    n_sigma,
+                    fg_thresh,
+                    seed_thresh,
+                    min_mask_sum,
+                    min_unclustered_sum,
+                    min_object_size,
+                    cluster,
+                    device,
+                )
+
+            base, _ = os.path.splitext(os.path.basename(sample["im_name"][0]))
             image_file_names.append(base)
 
-            if (one_hot):
-                if ('instance' in sample):
-                    all_results = obtain_APdsb_one_hot(gt_image=sample['instance'].squeeze().cpu().detach().numpy(),
-                                                       prediction_image=instance_map.cpu().detach().numpy(),
-                                                       ap_val=ap_val)
-                    if (verbose):
+            if one_hot:
+                if "instance" in sample:
+                    all_results = obtain_APdsb_one_hot(
+                        gt_image=sample["instance"].squeeze().cpu().detach().numpy(),
+                        prediction_image=instance_map.cpu().detach().numpy(),
+                        ap_val=ap_val,
+                    )
+                    if verbose:
                         print("Accuracy: {:.03f}".format(all_results), flush=True)
                     result_list.append(all_results)
             else:
-                if ('instance' in sample):
-                    all_results = matching_dataset(y_true=[sample['instance'].squeeze().cpu().detach().numpy()],
-                                                   y_pred=[instance_map.cpu().detach().numpy()], thresh=ap_val,
-                                                   show_progress=False)
-                    if (verbose):
-                        print("Accuracy: {:.03f}".format(all_results.accuracy), flush=True)
+                if "instance" in sample:
+                    all_results = matching_dataset(
+                        y_true=[sample["instance"].squeeze().cpu().detach().numpy()],
+                        y_pred=[instance_map.cpu().detach().numpy()],
+                        thresh=ap_val,
+                        show_progress=False,
+                    )
+                    if verbose:
+                        print(
+                            "Accuracy: {:.03f}".format(all_results.accuracy), flush=True
+                        )
                     result_list.append(all_results.accuracy)
 
             if save_images and ap_val == 0.5:
-
-                if not os.path.exists(os.path.join(save_dir, 'predictions/')):
-                    os.makedirs(os.path.join(save_dir, 'predictions/'))
-                    print("Created new directory {}".format(os.path.join(save_dir, 'predictions/')))
-                if not os.path.exists(os.path.join(save_dir, 'ground-truth/')):
-                    os.makedirs(os.path.join(save_dir, 'ground-truth/'))
-                    print("Created new directory {}".format(os.path.join(save_dir, 'ground-truth/')))
-                if not os.path.exists(os.path.join(save_dir, 'seeds/')):
-                    os.makedirs(os.path.join(save_dir, 'seeds/'))
-                    print("Created new directory {}".format(os.path.join(save_dir, 'seeds/')))
+                if not os.path.exists(os.path.join(save_dir, "predictions/")):
+                    os.makedirs(os.path.join(save_dir, "predictions/"))
+                    print(
+                        "Created new directory {}".format(
+                            os.path.join(save_dir, "predictions/")
+                        )
+                    )
+                if not os.path.exists(os.path.join(save_dir, "ground-truth/")):
+                    os.makedirs(os.path.join(save_dir, "ground-truth/"))
+                    print(
+                        "Created new directory {}".format(
+                            os.path.join(save_dir, "ground-truth/")
+                        )
+                    )
+                if not os.path.exists(os.path.join(save_dir, "seeds/")):
+                    os.makedirs(os.path.join(save_dir, "seeds/"))
+                    print(
+                        "Created new directory {}".format(
+                            os.path.join(save_dir, "seeds/")
+                        )
+                    )
 
                 # save predictions
-                base, _ = os.path.splitext(os.path.basename(sample['im_name'][0]))
-                instances_file = os.path.join(save_dir, 'predictions/', base + '.tif')
-                imsave(instances_file, instance_map.cpu().detach().numpy().astype(
-                    np.uint16))
-                seeds_file = os.path.join(save_dir, 'seeds/', base + '.tif')
+                base, _ = os.path.splitext(os.path.basename(sample["im_name"][0]))
+                instances_file = os.path.join(save_dir, "predictions/", base + ".tif")
+                imsave(
+                    instances_file,
+                    instance_map.cpu().detach().numpy().astype(np.uint16),
+                )
+                seeds_file = os.path.join(save_dir, "seeds/", base + ".tif")
                 imsave(seeds_file, seed_map.cpu().detach().numpy())
-                if ('instance' in sample):
-                    gt_file = os.path.join(save_dir, 'ground-truth/', base + '.tif')
-                    imsave(gt_file, sample['instance'].squeeze().cpu().detach().numpy())
-
-        if save_results and 'instance' in sample:
-            if not os.path.exists(os.path.join(save_dir, 'results/')):
-                os.makedirs(os.path.join(save_dir, 'results/'))
-                print("Created new directory {}".format(os.path.join(save_dir, 'results/')))
-            txt_file = os.path.join(save_dir,
-                                    'results/combined_AP-' + '{:.02f}'.format(ap_val) + '_tta-' + str(tta) + '.txt')
-            with open(txt_file, 'w') as f:
+                if "instance" in sample:
+                    gt_file = os.path.join(save_dir, "ground-truth/", base + ".tif")
+                    imsave(gt_file, sample["instance"].squeeze().cpu().detach().numpy())
+
+        if save_results and "instance" in sample:
+            if not os.path.exists(os.path.join(save_dir, "results/")):
+                os.makedirs(os.path.join(save_dir, "results/"))
+                print(
+                    "Created new directory {}".format(
+                        os.path.join(save_dir, "results/")
+                    )
+                )
+            txt_file = os.path.join(
+                save_dir,
+                "results/combined_AP-"
+                + "{:.02f}".format(ap_val)
+                + "_tta-"
+                + str(tta)
+                + ".txt",
+            )
+            with open(txt_file, "w") as f:
                 f.writelines(
-                    "image_file_name, min_mask_sum, min_unclustered_sum, min_object_size, seed_thresh, intersection_threshold, accuracy \n")
+                    "image_file_name, min_mask_sum, min_unclustered_sum, \
+                            min_object_size, seed_thresh, \
+                            intersection_threshold, accuracy \n"
+                )
                 f.writelines("+++++++++++++++++++++++++++++++++\n")
                 for ind, im_name in enumerate(image_file_names):
-                    im_name_tif = im_name + '.tif'
+                    im_name_tif = im_name + ".tif"
                     score = result_list[ind]
                     f.writelines(
-                        "{} {:.02f} {:.02f} {:.02f} {:.02f} {:.02f} {:.05f} \n".format(im_name_tif, min_mask_sum,
-                                                                                       min_unclustered_sum,
-                                                                                       min_object_size, seed_thresh,
-                                                                                       ap_val, score))
+                        "{} {:.02f} {:.02f} {:.02f} {:.02f} {:.02f} {:.05f} \n".format(
+                            im_name_tif,
+                            min_mask_sum,
+                            min_unclustered_sum,
+                            min_object_size,
+                            seed_thresh,
+                            ap_val,
+                            score,
+                        )
+                    )
                 f.writelines("+++++++++++++++++++++++++++++++++\n")
-                f.writelines("Average Precision (AP_dsb)  {:.02f} {:.05f}\n".format(ap_val, np.mean(result_list)))
+                f.writelines(
+                    "Average Precision (AP_dsb)  {:.02f} {:.05f}\n".format(
+                        ap_val, np.mean(result_list)
+                    )
+                )
 
         if len(result_list) != 0:
             print(
-                "Mean Average Precision (AP_dsb) at IOU threshold = {} at foreground threshold = {:.05f}, is equal to {:.05f}".format(
-                    ap_val, fg_thresh, np.mean(result_list)))
+                "Mean Average Precision (AP_dsb) at IOU threshold = {} \
+                        at foreground threshold = {:.05f}, \
+                        is equal to {:.05f}".format(
+                    ap_val, fg_thresh, np.mean(result_list)
+                )
+            )
             return -np.mean(result_list)
         else:
             return 0.0
@@ -761,33 +1208,33 @@ def test(fg_thresh, *args):
 def round_up_8(x):
     """Helper function for rounding integer to next multiple of 8
 
-        e.g:
-        round_up_8(10) = 16
-
-            Parameters
-            ----------
-            x : int
-                Integer
-            Returns
-            -------
-            int
-            """
+    e.g:
+    round_up_8(10) = 16
+
+        Parameters
+        ----------
+        x : int
+            Integer
+        Returns
+        -------
+        int
+    """
     return (int(x) + 7) & (-8)
 
 
 def pad_3d(im_tile):
     """Pad a 3D  image so that its dimensions are all multiples of 8
 
-                Parameters
-                ----------
-                im_tile : numpy array (D x H x W)
-                   3D Image which needs to be padded!
-                Returns
-                -------
-                (numpy array, int, int, int)
-                (Padded 3D image, diff in x, diff in y, diff in z)
-                The last three values are the amount of padding needed in the x, y and z dimensions
-                """
+    Parameters
+    ----------
+    im_tile : numpy array (D x H x W)
+       3D Image which needs to be padded!
+    Returns
+    -------
+    (numpy array, int, int, int)
+    (Padded 3D image, diff in x, diff in y, diff in z)
+    The last three values are the amount of padding needed in the x, y and z dimensions
+    """
 
     multiple_z = im_tile.shape[2] // 8
     multiple_y = im_tile.shape[3] // 8
@@ -807,37 +1254,64 @@ def pad_3d(im_tile):
 
     p3d = (0, diff_x, 0, diff_y, 0, diff_z)  # last dim, second last dim, third last dim
 
-    im_tile = F.pad(im_tile, p3d, 'reflect')
+    im_tile = F.pad(im_tile, p3d, "reflect")
     return im_tile, diff_x, diff_y, diff_z
 
 
 def test_3d(fg_thresh, *args):
     """Infer the trained 3D model on 3D images
 
-            Parameters
-            ----------
-            fg_thresh : float
-                foreground threshold decides which pixels are considered for clustering, based on the predicted seediness scores at these pixels.
-            args: dictionary
-                Contains other paremeters such as `ap_val`, `seed_thresh` etc
-            Returns
-            -------
-            float
-                Average `AP_dsb` over all test images
-            """
-    seed_thresh, ap_val, min_mask_sum, min_unclustered_sum, min_object_size, tta, model, dataset_it, \
-    save_images, save_results, save_dir, verbose, grid_x, grid_y, grid_z, \
-    pixel_x, pixel_y, pixel_z, one_hot, mask_region, n_sigma, cluster_fast, expand_grid, uniform_ds_factor = args
+    Parameters
+    ----------
+    fg_thresh : float
+        foreground threshold decides which pixels are considered for clustering,
+        based on the predicted seediness scores at these pixels.
+    args: dictionary
+        Contains other paremeters such as `ap_val`, `seed_thresh` etc
+    Returns
+    -------
+    float
+        Average `AP_dsb` over all test images
+    """
+    (
+        seed_thresh,
+        ap_val,
+        min_mask_sum,
+        min_unclustered_sum,
+        min_object_size,
+        tta,
+        model,
+        dataset_it,
+        save_images,
+        save_results,
+        save_dir,
+        verbose,
+        grid_x,
+        grid_y,
+        grid_z,
+        pixel_x,
+        pixel_y,
+        pixel_z,
+        one_hot,
+        mask_region,
+        n_sigma,
+        cluster_fast,
+        expand_grid,
+        uniform_ds_factor,
+        device,
+    ) = args
 
     model.eval()
     # cluster module
-    cluster = Cluster_3d(grid_z, grid_y, grid_x, pixel_z, pixel_y, pixel_x)
+    cluster = Cluster_3d(grid_z, grid_y, grid_x, pixel_z, pixel_y, pixel_x, device)
 
     with torch.no_grad():
         result_list = []
         image_file_names = []
         for sample in tqdm(dataset_it):
-            im = sample['image']
+            im = sample["image"]
+            im = im.to(device)
+
             D, H, W = im.shape[2], im.shape[3], im.shape[4]
 
             if D > grid_z or H > grid_y or W > grid_x:
@@ -848,109 +1322,216 @@ def test_3d(fg_thresh, *args):
                     W_ = temp
                     pixel_x_modified = pixel_y_modified = H_ / grid_y
                     pixel_z_modified = D_ * pixel_z / grid_z
-                    cluster = Cluster_3d(D_, H_, W_, pixel_z_modified, pixel_y_modified, pixel_x_modified)
-                    instance_map, seed_map = predict_3d(im, model, tta, cluster_fast,
-                                                     n_sigma, fg_thresh, seed_thresh, min_mask_sum, min_unclustered_sum,
-                                                     min_object_size, cluster)
+                    cluster = Cluster_3d(
+                        D_,
+                        H_,
+                        W_,
+                        pixel_z_modified,
+                        pixel_y_modified,
+                        pixel_x_modified,
+                        device,
+                    )
+                    instance_map, seed_map = predict_3d(
+                        im,
+                        model,
+                        tta,
+                        cluster_fast,
+                        n_sigma,
+                        fg_thresh,
+                        seed_thresh,
+                        min_mask_sum,
+                        min_unclustered_sum,
+                        min_object_size,
+                        cluster,
+                        device,
+                    )
                 else:
                     # here, we try stitching predictions instead
                     last = 1
                     instance_map = np.zeros((D, H, W), dtype=np.int16)
-                    seed_map = np.zeros((D, H, W), dtype=np.float)
+                    seed_map = np.zeros((D, H, W), dtype=float)
                     num_overlap_pixels = 4
-                    for z in range(0, D, grid_z -num_overlap_pixels):
+                    for z in range(0, D, grid_z - num_overlap_pixels):
                         for y in range(0, H, grid_y - num_overlap_pixels):
                             for x in range(0, W, grid_x - num_overlap_pixels):
-                                instance_map_tile, seed_map_tile = predict_3d(im[:, :, z:z + grid_z, y:y + grid_y, x:x + grid_x], model, tta,
-                                                                       cluster_fast,
-                                                                       n_sigma, fg_thresh, seed_thresh, min_mask_sum,
-                                                                       min_unclustered_sum, min_object_size,
-                                                                       cluster)
-                                last, instance_map = stitch_3d(instance_map_tile.cpu().detach().numpy(), instance_map, z, y, x,
-                                                           last, num_overlap_pixels)
-                                seed_map[z:z + grid_z, y:y + grid_y, x:x + grid_x] = seed_map_tile.cpu().detach().numpy()
-                    instance_map = torch.from_numpy(instance_map).cuda()
-                    seed_map = torch.from_numpy(seed_map).float().cuda()
+                                instance_map_tile, seed_map_tile = predict_3d(
+                                    im[
+                                        :,
+                                        :,
+                                        z : z + grid_z,
+                                        y : y + grid_y,
+                                        x : x + grid_x,
+                                    ],
+                                    model,
+                                    tta,
+                                    cluster_fast,
+                                    n_sigma,
+                                    fg_thresh,
+                                    seed_thresh,
+                                    min_mask_sum,
+                                    min_unclustered_sum,
+                                    min_object_size,
+                                    cluster,
+                                    device,
+                                )
+                                last, instance_map = stitch_3d(
+                                    instance_map_tile.cpu().detach().numpy(),
+                                    instance_map,
+                                    z,
+                                    y,
+                                    x,
+                                    last,
+                                    num_overlap_pixels,
+                                    device,
+                                )
+                                seed_map[
+                                    z : z + grid_z, y : y + grid_y, x : x + grid_x
+                                ] = (seed_map_tile.cpu().detach().numpy())
+                    instance_map = torch.from_numpy(instance_map).to(device)
+                    seed_map = torch.from_numpy(seed_map).float().to(device)
             else:
-                instance_map, seed_map = predict_3d(im, model, tta, cluster_fast,
-                                                 n_sigma, fg_thresh, seed_thresh, min_mask_sum, min_unclustered_sum,
-                                                 min_object_size, cluster)
-
+                instance_map, seed_map = predict_3d(
+                    im,
+                    model,
+                    tta,
+                    cluster_fast,
+                    n_sigma,
+                    fg_thresh,
+                    seed_thresh,
+                    min_mask_sum,
+                    min_unclustered_sum,
+                    min_object_size,
+                    cluster,
+                    device,
+                )
 
             # zoom back to original size
-            instance_map = zoom(instance_map.cpu().detach().numpy().astype(np.uint16), uniform_ds_factor, order=0)
+            instance_map = zoom(
+                instance_map.cpu().detach().numpy().astype(np.uint16),
+                uniform_ds_factor,
+                order=0,
+            )
             seed_map = zoom(seed_map.cpu().detach().numpy(), uniform_ds_factor)
 
-            if ('instance' in sample):
-                instances = sample['instance'].squeeze()  # Z, Y, X
-
-            if (mask_region is not None):
-                # ignore predictions in this region prior to saving the tiffs or prior to comparison with GT masks
-                instance_map[int(mask_region[0][0]):int(mask_region[1][0]),
-                int(mask_region[0][1]):int(mask_region[1][1]),
-                int(mask_region[0][2]):int(mask_region[1][2])] = 0  # Z Y X
+            if "instance" in sample:
+                instances = sample["instance"].squeeze()  # Z, Y, X
+
+            if mask_region is not None:
+                # ignore predictions in this region prior to saving the tiffs
+                # or prior to comparison with GT masks
+                instance_map[
+                    int(mask_region[0][0]) : int(mask_region[1][0]),
+                    int(mask_region[0][1]) : int(mask_region[1][1]),
+                    int(mask_region[0][2]) : int(mask_region[1][2]),
+                ] = 0  # Z Y X
             else:
                 pass
 
-            if ('instance' in sample):
-                instance_map = instance_map[:instances.shape[0], :instances.shape[1], :instances.shape[2]]
-                seed_map = seed_map[:instances.shape[0], :instances.shape[1], :instances.shape[2]]
-                all_results = matching_dataset(y_true=[instances.cpu().detach().numpy()],
-                                               y_pred=[instance_map], thresh=ap_val,
-                                               show_progress=False)
-                if (verbose):
+            if "instance" in sample:
+                instance_map = instance_map[
+                    : instances.shape[0], : instances.shape[1], : instances.shape[2]
+                ]
+                seed_map = seed_map[
+                    : instances.shape[0], : instances.shape[1], : instances.shape[2]
+                ]
+                all_results = matching_dataset(
+                    y_true=[instances.cpu().detach().numpy()],
+                    y_pred=[instance_map],
+                    thresh=ap_val,
+                    show_progress=False,
+                )
+                if verbose:
                     print("Accuracy: {:.03f}".format(all_results.accuracy), flush=True)
                 result_list.append(all_results.accuracy)
 
             if save_images and ap_val == 0.5:
-                if not os.path.exists(os.path.join(save_dir, 'predictions/')):
-                    os.makedirs(os.path.join(save_dir, 'predictions/'))
-                    print("Created new directory {}".format(os.path.join(save_dir, 'predictions/')))
-                if not os.path.exists(os.path.join(save_dir, 'ground-truth/')):
-                    os.makedirs(os.path.join(save_dir, 'ground-truth/'))
-                    print("Created new directory {}".format(os.path.join(save_dir, 'ground-truth/')))
-                if not os.path.exists(os.path.join(save_dir, 'seeds/')):
-                    os.makedirs(os.path.join(save_dir, 'seeds/'))
-                    print("Created new directory {}".format(os.path.join(save_dir, 'seeds/')))
-
-                base, _ = os.path.splitext(os.path.basename(sample['im_name'][0]))
+                if not os.path.exists(os.path.join(save_dir, "predictions/")):
+                    os.makedirs(os.path.join(save_dir, "predictions/"))
+                    print(
+                        "Created new directory {}".format(
+                            os.path.join(save_dir, "predictions/")
+                        )
+                    )
+                if not os.path.exists(os.path.join(save_dir, "ground-truth/")):
+                    os.makedirs(os.path.join(save_dir, "ground-truth/"))
+                    print(
+                        "Created new directory {}".format(
+                            os.path.join(save_dir, "ground-truth/")
+                        )
+                    )
+                if not os.path.exists(os.path.join(save_dir, "seeds/")):
+                    os.makedirs(os.path.join(save_dir, "seeds/"))
+                    print(
+                        "Created new directory {}".format(
+                            os.path.join(save_dir, "seeds/")
+                        )
+                    )
+
+                base, _ = os.path.splitext(os.path.basename(sample["im_name"][0]))
                 image_file_names.append(base)
 
-                instances_file = os.path.join(save_dir, 'predictions/', base + '.tif')
+                instances_file = os.path.join(save_dir, "predictions/", base + ".tif")
                 imsave(instances_file, instance_map)
 
-                seeds_file = os.path.join(save_dir, 'seeds/', base + '.tif')
+                seeds_file = os.path.join(save_dir, "seeds/", base + ".tif")
                 imsave(seeds_file, seed_map)
 
-                if ('instance' in sample):
-                    gt_file = os.path.join(save_dir, 'ground-truth/', base + '.tif')
-                    imsave(gt_file, sample['instance'].squeeze().cpu().detach().numpy())
-
-        if save_results and 'instance' in sample:
-            if not os.path.exists(os.path.join(save_dir, 'results/')):
-                os.makedirs(os.path.join(save_dir, 'results/'))
-                print("Created new directory {}".format(os.path.join(save_dir, 'results/')))
-            txt_file = os.path.join(save_dir,
-                                    'results/combined_AP-' + '{:.02f}'.format(ap_val) + '_tta-' + str(tta) + '.txt')
-            with open(txt_file, 'w') as f:
+                if "instance" in sample:
+                    gt_file = os.path.join(save_dir, "ground-truth/", base + ".tif")
+                    imsave(gt_file, sample["instance"].squeeze().cpu().detach().numpy())
+
+        if save_results and "instance" in sample:
+            if not os.path.exists(os.path.join(save_dir, "results/")):
+                os.makedirs(os.path.join(save_dir, "results/"))
+                print(
+                    "Created new directory {}".format(
+                        os.path.join(save_dir, "results/")
+                    )
+                )
+            txt_file = os.path.join(
+                save_dir,
+                "results/combined_AP-"
+                + "{:.02f}".format(ap_val)
+                + "_tta-"
+                + str(tta)
+                + ".txt",
+            )
+            with open(txt_file, "w") as f:
                 f.writelines(
-                    "image_file_name, min_mask_sum, min_unclustered_sum, min_object_size, seed_thresh, intersection_threshold, accuracy \n")
+                    "image_file_name, min_mask_sum, min_unclustered_sum, \
+                            min_object_size, seed_thresh, \
+                            intersection_threshold, accuracy \n"
+                )
                 f.writelines("+++++++++++++++++++++++++++++++++\n")
                 for ind, im_name in enumerate(image_file_names):
-                    im_name_tif = im_name + '.tif'
+                    im_name_tif = im_name + ".tif"
                     score = result_list[ind]
                     f.writelines(
-                        "{} {:.02f} {:.02f} {:.02f} {:.02f} {:.02f} {:.05f} \n".format(im_name_tif, min_mask_sum,
-                                                                                       min_unclustered_sum,
-                                                                                       min_object_size, seed_thresh,
-                                                                                       ap_val, score))
+                        "{} {:.02f} {:.02f} {:.02f} {:.02f} {:.02f} {:.05f} \n".format(
+                            im_name_tif,
+                            min_mask_sum,
+                            min_unclustered_sum,
+                            min_object_size,
+                            seed_thresh,
+                            ap_val,
+                            score,
+                        )
+                    )
                 f.writelines("+++++++++++++++++++++++++++++++++\n")
-                f.writelines("Average Precision (AP_dsb)  {:.02f} {:.05f}\n".format(ap_val, np.mean(result_list)))
+                f.writelines(
+                    "Average Precision (AP_dsb)  {:.02f} {:.05f}\n".format(
+                        ap_val, np.mean(result_list)
+                    )
+                )
 
         if len(result_list) != 0:
             print(
-                "Mean Average Precision (AP_dsb) at IOU threshold = {} at seediness threshold = {:.05f}, is equal to {:.05f}".format(
-                    ap_val, seed_thresh, np.mean(result_list)))
+                "Mean Average Precision (AP_dsb) at IOU threshold \
+                        = {} at seediness threshold = {:.05f}, is equal to \
+                        {:.05f}".format(
+                    ap_val, seed_thresh, np.mean(result_list)
+                )
+            )
             return -np.mean(result_list)
         else:
             return 0.0
@@ -959,30 +1540,54 @@ def test_3d(fg_thresh, *args):
 def test_3d_sliced(fg_thresh, *args):
     """Infer the trained 2D model on 3D images in a `sliced` fashion
 
-                Parameters
-                ----------
-                fg_thresh : float
-                    foreground threshold decides which pixels are considered for clustering, based on the predicted seediness scores at these pixels.
-                args: dictionary
-                    Contains other paremeters such as `ap_val`, `seed_thresh` etc
-                Returns
-                -------
-                float
-                    Average `AP_dsb` over all test images
-                """
-    seed_thresh, ap_val, min_mask_sum, min_unclustered_sum, min_object_size, tta, model, dataset_it, save_images, save_results, save_dir, \
-    verbose, grid_x, grid_y, grid_z, pixel_x, pixel_y, pixel_z, \
-    one_hot, mask_region, n_sigma, anisotropy_factor = args
+    Parameters
+    ----------
+    fg_thresh : float
+        foreground threshold decides which pixels are considered
+        for clustering, based on the predicted seediness scores at these pixels.
+    args: dictionary
+        Contains other paremeters such as `ap_val`, `seed_thresh` etc
+    Returns
+    -------
+    float
+        Average `AP_dsb` over all test images
+    """
+    (
+        seed_thresh,
+        ap_val,
+        min_mask_sum,
+        min_unclustered_sum,
+        min_object_size,
+        tta,
+        model,
+        dataset_it,
+        save_images,
+        save_results,
+        save_dir,
+        verbose,
+        grid_x,
+        grid_y,
+        grid_z,
+        pixel_x,
+        pixel_y,
+        pixel_z,
+        one_hot,
+        mask_region,
+        n_sigma,
+        anisotropy_factor,
+        device,
+    ) = args
 
     model.eval()
     # cluster module
-    cluster = Cluster_3d(grid_z, grid_y, grid_x, pixel_z, pixel_y, pixel_x)
+    cluster = Cluster_3d(grid_z, grid_y, grid_x, pixel_z, pixel_y, pixel_x, device)
 
     with torch.no_grad():
         result_list = []
         image_file_names = []
         for sample in tqdm(dataset_it):
-            im = sample['image']  # isotropically expanded image
+            im = sample["image"]  # isotropically expanded image
+            im = im.to(device)
             multiple_z = im.shape[2] // 8
             multiple_y = im.shape[3] // 8
             multiple_x = im.shape[4] // 8
@@ -999,27 +1604,39 @@ def test_3d_sliced(fg_thresh, *args):
                 diff_x = 8 * (multiple_x + 1) - im.shape[4]
             else:
                 diff_x = 0
-            p3d = (diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2, diff_z // 2,
-                   diff_z - diff_z // 2)  # last dim, second last dim, third last dim!
+            p3d = (
+                diff_x // 2,
+                diff_x - diff_x // 2,
+                diff_y // 2,
+                diff_y - diff_y // 2,
+                diff_z // 2,
+                diff_z - diff_z // 2,
+            )  # last dim, second last dim, third last dim!
 
             im = F.pad(im, p3d, "reflect")
-            if ('instance' in sample):
-                instances = sample['instance'].squeeze()  # isotropically expanded GT instance map
+            if "instance" in sample:
+                instances = sample[
+                    "instance"
+                ].squeeze()  # isotropically expanded GT instance map
 
             print("Processing `YX` slices .....")
             # obtain YX slices
-            output_YX_3d = torch.zeros(1, 5, im.shape[2], im.shape[3], im.shape[4]).cuda()
+            output_YX_3d = torch.zeros(1, 5, im.shape[2], im.shape[3], im.shape[4]).to(
+                device
+            )
             for z in tqdm(range(im.shape[2]), position=0, leave=True):
-                if (tta):
+                if tta:
                     output_YX = apply_tta_2d(im[:, :, z, ...], model)
                 else:
                     output_YX = model(im[:, :, z, ...])
                 output_YX_3d[:, :, z, ...] = output_YX
             print("Processing `ZX` slices .....")
             # obtain ZX slices
-            output_ZX_3d = torch.zeros(1, 5, im.shape[2], im.shape[3], im.shape[4]).cuda()
+            output_ZX_3d = torch.zeros(1, 5, im.shape[2], im.shape[3], im.shape[4]).to(
+                device
+            )
             for y in tqdm(range(im.shape[3]), position=0, leave=True):
-                if (tta):
+                if tta:
                     output_ZX = apply_tta_2d(im[:, :, :, y, ...], model)
                 else:
                     output_ZX = model(im[:, :, :, y, ...])
@@ -1027,130 +1644,211 @@ def test_3d_sliced(fg_thresh, *args):
                 output_ZX_3d[:, :, :, y, ...] = output_ZX
             print("Processing `ZY` slices .....")
             # obtain ZY slices
-            output_ZY_3d = torch.zeros(1, 5, im.shape[2], im.shape[3], im.shape[4]).cuda()
+            output_ZY_3d = torch.zeros(1, 5, im.shape[2], im.shape[3], im.shape[4]).to(
+                device
+            )
             for x in tqdm(range(im.shape[4]), position=0, leave=True):
-                if (tta):
+                if tta:
                     output_ZY = apply_tta_2d(im[:, :, :, :, x], model)
                 else:
                     output_ZY = model(im[:, :, :, :, x])
 
                 # offset in y, offset in z, sigma in y, sigma in z, seed
                 output_ZY_3d[:, :, :, :, x] = output_ZY
-            output = torch.from_numpy(
-                np.zeros((1, 7, im.shape[2], im.shape[3], im.shape[4]))).float().cuda()  # B (6+1) Z Y X
+            output = (
+                torch.from_numpy(
+                    np.zeros((1, 7, im.shape[2], im.shape[3], im.shape[4]))
+                )
+                .float()
+                .to(device)
+            )  # B (6+1) Z Y X
 
             # handle seediness
-            output[:, 6, ...] = 1 / 3 * (output_YX_3d[:, 4, ...] + output_ZX_3d[:, 4, ...] + output_ZY_3d[:, 4, ...])
+            output[:, 6, ...] = (
+                1
+                / 3
+                * (
+                    output_YX_3d[:, 4, ...]
+                    + output_ZX_3d[:, 4, ...]
+                    + output_ZY_3d[:, 4, ...]
+                )
+            )
 
             # handle offset in X
-            output[:, 0, ...] = 1 / 2 * (output_YX_3d[:, 0, ...] + output_ZX_3d[:, 0, ...])
+            output[:, 0, ...] = (
+                1 / 2 * (output_YX_3d[:, 0, ...] + output_ZX_3d[:, 0, ...])
+            )
 
             # handle offset in Y
-            output[:, 1, ...] = 1 / 2 * (output_YX_3d[:, 1, ...] + output_ZY_3d[:, 0, ...])
+            output[:, 1, ...] = (
+                1 / 2 * (output_YX_3d[:, 1, ...] + output_ZY_3d[:, 0, ...])
+            )
 
             # handle offset in Z
-            output[:, 2, ...] = 1 / 2 * (output_ZX_3d[:, 1, ...] + output_ZY_3d[:, 1, ...])
+            output[:, 2, ...] = (
+                1 / 2 * (output_ZX_3d[:, 1, ...] + output_ZY_3d[:, 1, ...])
+            )
 
             # sigma in X
-            output[:, 3, ...] = 1 / 2 * (output_YX_3d[:, 2, ...] + output_ZX_3d[:, 2, ...])
+            output[:, 3, ...] = (
+                1 / 2 * (output_YX_3d[:, 2, ...] + output_ZX_3d[:, 2, ...])
+            )
 
             # sigma in Y
-            output[:, 4, ...] = 1 / 2 * (output_YX_3d[:, 3, ...] + output_ZY_3d[:, 2, ...])
+            output[:, 4, ...] = (
+                1 / 2 * (output_YX_3d[:, 3, ...] + output_ZY_3d[:, 2, ...])
+            )
 
             # sigma in Z
-            output[:, 5, ...] = 1 / 2 * (output_ZX_3d[:, 3, ...] + output_ZY_3d[:, 3, ...])
+            output[:, 5, ...] = (
+                1 / 2 * (output_ZX_3d[:, 3, ...] + output_ZY_3d[:, 3, ...])
+            )
 
             # unpad output, instances and image
-            if (diff_z - diff_z // 2) is not 0:
-                output = output[:, :, diff_z // 2:-(diff_z - diff_z // 2), ...]
+            if (diff_z - diff_z // 2) != 0:
+                output = output[:, :, diff_z // 2 : -(diff_z - diff_z // 2), ...]
 
-            if (diff_y - diff_y // 2) is not 0:
-                output = output[:, :, :, diff_y // 2:-(diff_y - diff_y // 2), ...]
+            if (diff_y - diff_y // 2) != 0:
+                output = output[:, :, :, diff_y // 2 : -(diff_y - diff_y // 2), ...]
 
-            if (diff_x - diff_x // 2) is not 0:
-                output = output[..., diff_x // 2:-(diff_x - diff_x // 2)]
+            if (diff_x - diff_x // 2) != 0:
+                output = output[..., diff_x // 2 : -(diff_x - diff_x // 2)]
 
             # reverse the isotropic sampling
             output_cpu = output.cpu().detach().numpy()  # BCZYX
             output_cpu = zoom(output_cpu, (1, 1, 1 / anisotropy_factor, 1, 1), order=0)
-            output = torch.from_numpy(output_cpu).float().cuda()
-            if ('instance' in sample):
-                gt_image = zoom(instances.cpu().detach().numpy(), (1 / anisotropy_factor, 1, 1), order=0)
-
-            instance_map = cluster.cluster(output[0],
-                                           n_sigma=n_sigma,  # 3
-                                           seed_thresh=seed_thresh,
-                                           fg_thresh=fg_thresh,
-                                           min_mask_sum=min_mask_sum,
-                                           min_unclustered_sum=min_unclustered_sum,
-                                           min_object_size=min_object_size,
-                                           )
-
-            if (mask_region is not None):
-                # ignore predictions in this region prior to saving the tiffs or prior to comparison with GT masks
-                instance_map[int(mask_region[0][0]):int(mask_region[1][0]),
-                int(mask_region[0][1]):int(mask_region[1][1]),
-                int(mask_region[0][2]):int(mask_region[1][2])] = 0  # Z Y X
+            output = torch.from_numpy(output_cpu).float().to(device)
+            if "instance" in sample:
+                zoom(
+                    instances.cpu().detach().numpy(),
+                    (1 / anisotropy_factor, 1, 1),
+                    order=0,
+                )
+
+            instance_map = cluster.cluster(
+                output[0],
+                n_sigma=n_sigma,  # 3
+                seed_thresh=seed_thresh,
+                fg_thresh=fg_thresh,
+                min_mask_sum=min_mask_sum,
+                min_unclustered_sum=min_unclustered_sum,
+                min_object_size=min_object_size,
+            )
+
+            if mask_region is not None:
+                # ignore predictions in this region prior to saving
+                # the tiffs or prior to comparison with GT masks
+                instance_map[
+                    int(mask_region[0][0]) : int(mask_region[1][0]),
+                    int(mask_region[0][1]) : int(mask_region[1][1]),
+                    int(mask_region[0][2]) : int(mask_region[1][2]),
+                ] = 0  # Z Y X
             else:
                 pass
 
-            if ('instance' in sample):
-                all_results = matching_dataset(y_true=[sample['instance'].squeeze().cpu().detach().numpy()],
-                                               y_pred=[instance_map.cpu().detach().numpy()],
-                                               thresh=ap_val,
-                                               show_progress=False)
-                if (verbose):
+            if "instance" in sample:
+                all_results = matching_dataset(
+                    y_true=[sample["instance"].squeeze().cpu().detach().numpy()],
+                    y_pred=[instance_map.cpu().detach().numpy()],
+                    thresh=ap_val,
+                    show_progress=False,
+                )
+                if verbose:
                     print("Accuracy: {:.03f}".format(all_results.accuracy), flush=True)
                 result_list.append(all_results.accuracy)
 
             if save_images and ap_val == 0.5:
-                if not os.path.exists(os.path.join(save_dir, 'predictions/')):
-                    os.makedirs(os.path.join(save_dir, 'predictions/'))
-                    print("Created new directory {}".format(os.path.join(save_dir, 'predictions/')))
-                if not os.path.exists(os.path.join(save_dir, 'seeds/')):
-                    os.makedirs(os.path.join(save_dir, 'seeds/'))
-                    print("Created new directory {}".format(os.path.join(save_dir, 'seeds/')))
-                if not os.path.exists(os.path.join(save_dir, 'ground-truth/')):
-                    os.makedirs(os.path.join(save_dir, 'ground-truth/'))
-                    print("Created new directory {}".format(os.path.join(save_dir, 'ground-truth/')))
-                base, _ = os.path.splitext(os.path.basename(sample['im_name'][0]))
+                if not os.path.exists(os.path.join(save_dir, "predictions/")):
+                    os.makedirs(os.path.join(save_dir, "predictions/"))
+                    print(
+                        "Created new directory {}".format(
+                            os.path.join(save_dir, "predictions/")
+                        )
+                    )
+                if not os.path.exists(os.path.join(save_dir, "seeds/")):
+                    os.makedirs(os.path.join(save_dir, "seeds/"))
+                    print(
+                        "Created new directory {}".format(
+                            os.path.join(save_dir, "seeds/")
+                        )
+                    )
+                if not os.path.exists(os.path.join(save_dir, "ground-truth/")):
+                    os.makedirs(os.path.join(save_dir, "ground-truth/"))
+                    print(
+                        "Created new directory {}".format(
+                            os.path.join(save_dir, "ground-truth/")
+                        )
+                    )
+                base, _ = os.path.splitext(os.path.basename(sample["im_name"][0]))
                 image_file_names.append(base)
 
-                prediction_file = os.path.join(save_dir, 'predictions/', base + '.tif')
-                imsave(prediction_file, instance_map.cpu().detach().numpy().astype(np.uint16))
-
-                seeds_file = os.path.join(save_dir, 'seeds/', base + '.tif')
-                imsave(seeds_file, torch.sigmoid(output[0, -1, ...]).cpu().detach().numpy())
-
-                if ('instance' in sample):
-                    gt_file = os.path.join(save_dir, 'ground-truth/', base + '.tif')
-                    imsave(gt_file, sample['instance'].squeeze().cpu().detach().numpy())
-
-        if save_results and 'instance' in sample:
-            if not os.path.exists(os.path.join(save_dir, 'results/')):
-                os.makedirs(os.path.join(save_dir, 'results/'))
-                print("Created new directory {}".format(os.path.join(save_dir, 'results/')))
-            txt_file = os.path.join(save_dir,
-                                    'results/combined_AP-' + '{:.02f}'.format(ap_val) + '_tta-' + str(tta) + '.txt')
-            with open(txt_file, 'w') as f:
+                prediction_file = os.path.join(save_dir, "predictions/", base + ".tif")
+                imsave(
+                    prediction_file,
+                    instance_map.cpu().detach().numpy().astype(np.uint16),
+                )
+
+                seeds_file = os.path.join(save_dir, "seeds/", base + ".tif")
+                imsave(
+                    seeds_file, torch.sigmoid(output[0, -1, ...]).cpu().detach().numpy()
+                )
+
+                if "instance" in sample:
+                    gt_file = os.path.join(save_dir, "ground-truth/", base + ".tif")
+                    imsave(gt_file, sample["instance"].squeeze().cpu().detach().numpy())
+
+        if save_results and "instance" in sample:
+            if not os.path.exists(os.path.join(save_dir, "results/")):
+                os.makedirs(os.path.join(save_dir, "results/"))
+                print(
+                    "Created new directory {}".format(
+                        os.path.join(save_dir, "results/")
+                    )
+                )
+            txt_file = os.path.join(
+                save_dir,
+                "results/combined_AP-"
+                + "{:.02f}".format(ap_val)
+                + "_tta-"
+                + str(tta)
+                + ".txt",
+            )
+            with open(txt_file, "w") as f:
                 f.writelines(
-                    "image_file_name, min_mask_sum, min_unclustered_sum, min_object_size, seed_thresh, intersection_threshold, accuracy \n")
+                    "image_file_name, min_mask_sum, min_unclustered_sum, \
+                            min_object_size, seed_thresh, \
+                            intersection_threshold, accuracy \n"
+                )
                 f.writelines("+++++++++++++++++++++++++++++++++\n")
                 for ind, im_name in enumerate(image_file_names):
-                    im_name_tif = im_name + '.tif'
+                    im_name_tif = im_name + ".tif"
                     score = result_list[ind]
                     f.writelines(
-                        "{} {:.02f} {:.02f} {:.02f} {:.02f} {:.02f} {:.05f} \n".format(im_name_tif, min_mask_sum,
-                                                                                       min_unclustered_sum,
-                                                                                       min_object_size, seed_thresh,
-                                                                                       ap_val, score))
+                        "{} {:.02f} {:.02f} {:.02f} {:.02f} {:.02f} {:.05f} \n".format(
+                            im_name_tif,
+                            min_mask_sum,
+                            min_unclustered_sum,
+                            min_object_size,
+                            seed_thresh,
+                            ap_val,
+                            score,
+                        )
+                    )
                 f.writelines("+++++++++++++++++++++++++++++++++\n")
-                f.writelines("Average Precision (AP_dsb)  {:.02f} {:.05f}\n".format(ap_val, np.mean(result_list)))
+                f.writelines(
+                    "Average Precision (AP_dsb)  {:.02f} {:.05f}\n".format(
+                        ap_val, np.mean(result_list)
+                    )
+                )
 
         if len(result_list) != 0:
             print(
-                "Mean Average Precision (AP_dsb) at IOU threshold = {} at seediness threshold = {:.05f}, is equal to {:.05f}".format(
-                    ap_val, seed_thresh, np.mean(result_list)))
+                "Mean Average Precision (AP_dsb) at IOU threshold \
+                        = {} at seediness threshold = {:.05f}, is equal to \
+                        {:.05f}".format(
+                    ap_val, seed_thresh, np.mean(result_list)
+                )
+            )
             return -np.mean(result_list)
         else:
             return 0.0
@@ -1159,22 +1857,25 @@ def test_3d_sliced(fg_thresh, *args):
 def perform_ilp(instance_map_z, min_depth=3, mean_object_size=3363):
     """Perform integer linear programming to minimize a cost function
 
-                Parameters
-                ----------
-                instance_map_z : list
-                                List of instance segmentations (numpy arrays) for all z slices
-                min_depth: int, optional
-                    Allows building edges to slices which are `min_depth` slices away.
-                    This is useful in case for one of the slices, an object is not segmented.
-                mean_object_size: int, optional
-                    The mean object size as calculated from the GT masks available for the training data.
-                    This is useful to compute the cost of appearance and disappearance.
-                    For example, a large segmented object should have a high cost appearance and disappearance cost.
-                Returns
-                -------
-                numpy array (D x H x W)
-                    The instance segmentation for the entire 3D image
-                """
+    Parameters
+    ----------
+    instance_map_z : list
+                    List of instance segmentations (numpy arrays) for all z slices
+    min_depth: int, optional
+        Allows building edges to slices
+        which are `min_depth` slices away.
+        This is useful in case for one of the slices, an object != segmented.
+    mean_object_size: int, optional
+        The mean object size as calculated from the GT masks available
+        for the training data.
+        This is useful to compute the cost of appearance and disappearance.
+        For example, a large segmented object should have a high cost
+        appearance and disappearance cost.
+    Returns
+    -------
+    numpy array (D x H x W)
+        The instance segmentation for the entire 3D image
+    """
 
     import gurobipy as gp
     import networkx as nx
@@ -1207,18 +1908,25 @@ def calculate_iou(pred, label):
         X_nodes_disappearance[z1] = {}
         for id in ids1:
             y, x = np.where(instance_map[z1] == id)
-            G.add_node(str(z1) + '_' + str(id), cost_node=1.0,
-                       cost_appearance=len(y) / mean_object_size,
-                       cost_disappearance=len(y) / mean_object_size)  # we just set it to a constant value for now
+            G.add_node(
+                str(z1) + "_" + str(id),
+                cost_node=1.0,
+                cost_appearance=len(y) / mean_object_size,
+                cost_disappearance=len(y) / mean_object_size,
+            )  # we just set it to a constant value for now
             if z1 == 0:
                 X_nodes_appearance[z1][id] = 1.0
             else:
-                X_nodes_appearance[z1][id] = model.addVar(vtype=GRB.BINARY, name="Xapp_%d_%d" % (z1, id))
+                X_nodes_appearance[z1][id] = model.addVar(
+                    vtype=GRB.BINARY, name="Xapp_%d_%d" % (z1, id)
+                )
 
             if z1 == instance_map.shape[0] - 1:
                 X_nodes_disappearance[z1][id] = 1.0
             else:
-                X_nodes_disappearance[z1][id] = model.addVar(vtype=GRB.BINARY, name="Xdisapp_%d_%d" % (z1, id))
+                X_nodes_disappearance[z1][id] = model.addVar(
+                    vtype=GRB.BINARY, name="Xdisapp_%d_%d" % (z1, id)
+                )
 
     print("Nodes added ...", flush=True)
 
@@ -1234,11 +1942,18 @@ def calculate_iou(pred, label):
             for id1 in ids1:
                 X_edges[z1][z2][id1] = {}
                 for id2 in ids2:
-                    iou = calculate_iou(instance_map[z1] == id1, instance_map[z2] == id2)
+                    iou = calculate_iou(
+                        instance_map[z1] == id1, instance_map[z2] == id2
+                    )
                     if iou > 0:
-                        G.add_edge(str(z1) + '_' + str(id1), str(z2) + '_' + str(id2), cost_edge=1 - iou / (z2 - z1))
-                        X_edges[z1][z2][id1][id2] = model.addVar(vtype=GRB.BINARY,
-                                                                 name="X_%d_%d_%d_%d" % (z1, id1, z2, id2))
+                        G.add_edge(
+                            str(z1) + "_" + str(id1),
+                            str(z2) + "_" + str(id2),
+                            cost_edge=1 - iou / (z2 - z1),
+                        )
+                        X_edges[z1][z2][id1][id2] = model.addVar(
+                            vtype=GRB.BINARY, name="X_%d_%d_%d_%d" % (z1, id1, z2, id2)
+                        )
     print("Edges added ...", flush=True)
 
     model.update()
@@ -1265,8 +1980,11 @@ def calculate_iou(pred, label):
             else:
                 incoming_indicators += X_edges[z2][z1][id2][id1]
         constraints.append(
-            model.addConstr(1.0 <= X_nodes_appearance[z1][id1] + incoming_indicators,
-                            'constraint_continuation_incoming_%d_%d' % (z1, id1)))  # TODO
+            model.addConstr(
+                1.0 <= X_nodes_appearance[z1][id1] + incoming_indicators,
+                "constraint_continuation_incoming_%d_%d" % (z1, id1),
+            )
+        )  # TODO
 
         outgoing_indicators = 0
         for edge in G.edges(node):
@@ -1278,8 +1996,11 @@ def calculate_iou(pred, label):
             else:
                 pass
         constraints.append(
-            model.addConstr(1.0 <= X_nodes_disappearance[z1][id1] + outgoing_indicators,
-                            'constraint_continuation_outgoing_%d_%d' % (z1, id1)))  # TODO
+            model.addConstr(
+                1.0 <= X_nodes_disappearance[z1][id1] + outgoing_indicators,
+                "constraint_continuation_outgoing_%d_%d" % (z1, id1),
+            )
+        )  # TODO
     print("Constraints added ...", flush=True)
 
     # finally consider costs
@@ -1288,15 +2009,17 @@ def calculate_iou(pred, label):
         # u = z_id
         z1, id1 = u.split("_")
         z1, id1 = int(z1), int(id1)
-        objective += cost_dic['cost_appearance'] * X_nodes_appearance[z1][id1] + cost_dic['cost_disappearance'] * \
-                     X_nodes_disappearance[z1][id1]  # TODO
+        objective += (
+            cost_dic["cost_appearance"] * X_nodes_appearance[z1][id1]
+            + cost_dic["cost_disappearance"] * X_nodes_disappearance[z1][id1]
+        )  # TODO
 
     for u, v, cost_dic in G.edges(data=True):
         # u = z_id
         z1, id1 = u.split("_")
         z2, id2 = v.split("_")
         z1, id1, z2, id2 = int(z1), int(id1), int(z2), int(id2)
-        objective += cost_dic['cost_edge'] * X_edges[z1][z2][id1][id2]
+        objective += cost_dic["cost_edge"] * X_edges[z1][z2][id1][id2]
 
     model.setObjective(objective)
     model.update()
@@ -1308,55 +2031,87 @@ def calculate_iou(pred, label):
     count = 1
     new_ids_dic = {}
     for v in tqdm(model.getVars(), position=0, leave=True):
-        if v.VarName.split('_')[0] == 'X':
-            _, z1, id1, z2, id2 = v.VarName.split('_')
+        if v.VarName.split("_")[0] == "X":
+            _, z1, id1, z2, id2 = v.VarName.split("_")
             indicator = v.X
             z1, id1, z2, id2 = int(z1), int(id1), int(z2), int(id2)
             if indicator:
-                if str(z1) + '_' + str(id1) in new_ids_dic.keys():
-                    instance_map_reassign[z1][instance_map[z1] == id1] = new_ids_dic[str(z1) + '_' + str(id1)]
-                    instance_map_reassign[z2][instance_map[z2] == id2] = new_ids_dic[str(z1) + '_' + str(id1)]
-                    new_ids_dic[str(z1) + '_' + str(id1)] = new_ids_dic[str(z1) + '_' + str(id1)]
-                    new_ids_dic[str(z2) + '_' + str(id2)] = new_ids_dic[str(z1) + '_' + str(id1)]
+                if str(z1) + "_" + str(id1) in new_ids_dic.keys():
+                    instance_map_reassign[z1][instance_map[z1] == id1] = new_ids_dic[
+                        str(z1) + "_" + str(id1)
+                    ]
+                    instance_map_reassign[z2][instance_map[z2] == id2] = new_ids_dic[
+                        str(z1) + "_" + str(id1)
+                    ]
+                    new_ids_dic[str(z1) + "_" + str(id1)] = new_ids_dic[
+                        str(z1) + "_" + str(id1)
+                    ]
+                    new_ids_dic[str(z2) + "_" + str(id2)] = new_ids_dic[
+                        str(z1) + "_" + str(id1)
+                    ]
                 else:
                     # start a new id
                     instance_map_reassign[z1][instance_map[z1] == id1] = count
                     instance_map_reassign[z2][instance_map[z2] == id2] = count
-                    new_ids_dic[str(z1) + '_' + str(id1)] = count
-                    new_ids_dic[str(z2) + '_' + str(id2)] = count
+                    new_ids_dic[str(z1) + "_" + str(id1)] = count
+                    new_ids_dic[str(z2) + "_" + str(id2)] = count
                     count += 1
     print("Ids reassigned  ...", flush=True)
     return instance_map_reassign.astype(np.uint16)
 
 
 def test_3d_ilp(fg_thresh, *args):
-    """Infer the trained 2D model on 3D images and then stitch the independent predictions on 2D z-slices by solving a global optimization task
-
-                Parameters
-                ----------
-                fg_thresh : float
-                    foreground threshold decides which pixels are considered for clustering, based on the predicted seediness scores at these pixels.
-                args: dictionary
-                    Contains other paremeters such as `ap_val`, `seed_thresh` etc
-                Returns
-                -------
-                float
-                    Average `AP_dsb` over all test images
-                """
-    seed_thresh, ap_val, min_mask_sum, min_unclustered_sum, min_object_size, mean_object_size, tta, model, dataset_it, save_images, save_results, save_dir, \
-    verbose, grid_x, grid_y, grid_z, pixel_x, pixel_y, pixel_z, \
-    one_hot, mask_region, n_sigma, anisotropy_factor = args
+    """Infer the trained 2D model on 3D images and then stitch the
+    independent predictions on 2D z-slices by solving a global optimization task
+
+    Parameters
+    ----------
+    fg_thresh : float
+        foreground threshold decides which pixels are considered for
+        clustering, based on the predicted seediness scores at these pixels.
+    args: dictionary
+        Contains other paremeters such as `ap_val`, `seed_thresh` etc
+    Returns
+    -------
+    float
+        Average `AP_dsb` over all test images
+    """
+    (
+        seed_thresh,
+        ap_val,
+        min_mask_sum,
+        min_unclustered_sum,
+        min_object_size,
+        mean_object_size,
+        tta,
+        model,
+        dataset_it,
+        save_images,
+        save_results,
+        save_dir,
+        verbose,
+        grid_x,
+        grid_y,
+        grid_z,
+        pixel_x,
+        pixel_y,
+        pixel_z,
+        one_hot,
+        mask_region,
+        n_sigma,
+        anisotropy_factor,
+        device,
+    ) = args
 
     model.eval()
     # cluster module
-    cluster = Cluster(grid_y, grid_x, pixel_y, pixel_x)
+    cluster = Cluster(grid_y, grid_x, pixel_y, pixel_x, device)
 
     with torch.no_grad():
         result_list = []
         image_file_names = []
         for sample in tqdm(dataset_it):
-
-            im = sample['image']  # B 1 Z Y X
+            im = sample["image"]  # B 1 Z Y X
             instance_map_z = []
             seed_z = []
             for z in tqdm(range(im.shape[2]), position=0, leave=True):
@@ -1372,90 +2127,140 @@ def test_3d_ilp(fg_thresh, *args):
                 else:
                     diff_x = 0
                 p2d = (
-                    diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2)  # last dim, second last dim
+                    diff_x // 2,
+                    diff_x - diff_x // 2,
+                    diff_y // 2,
+                    diff_y - diff_y // 2,
+                )  # last dim, second last dim
 
                 im_z = F.pad(im[:, :, z], p2d, "reflect")  # B C Y X
-                if (tta):
+                if tta:
                     output = apply_tta_2d(im_z, model)
                 else:
                     output = model(im_z)
 
-                instance_map, predictions = cluster.cluster(output[0],
-                                                            n_sigma=n_sigma,
-                                                            fg_thresh=fg_thresh,
-                                                            seed_thresh=seed_thresh,
-                                                            min_mask_sum=min_mask_sum,
-                                                            min_unclustered_sum=min_unclustered_sum,
-                                                            min_object_size=min_object_size)
+                instance_map, predictions = cluster.cluster(
+                    output[0],
+                    n_sigma=n_sigma,
+                    fg_thresh=fg_thresh,
+                    seed_thresh=seed_thresh,
+                    min_mask_sum=min_mask_sum,
+                    min_unclustered_sum=min_unclustered_sum,
+                    min_object_size=min_object_size,
+                )
 
                 # unpad instance_map, instances and images
-                if (diff_y - diff_y // 2) is not 0:
-                    instance_map = instance_map[diff_y // 2:-(diff_y - diff_y // 2), ...]
-                if (diff_x - diff_x // 2) is not 0:
-                    instance_map = instance_map[..., diff_x // 2:-(diff_x - diff_x // 2)]
+                if (diff_y - diff_y // 2) != 0:
+                    instance_map = instance_map[
+                        diff_y // 2 : -(diff_y - diff_y // 2), ...
+                    ]
+                if (diff_x - diff_x // 2) != 0:
+                    instance_map = instance_map[
+                        ..., diff_x // 2 : -(diff_x - diff_x // 2)
+                    ]
                 instance_map_z.append(instance_map.cpu().detach().numpy())
                 seed_z.append(torch.sigmoid(output[0, -1, ...]).cpu().detach().numpy())
             # correct predictions with gurobi
-            # imsave(save_dir+'/temp_pred.tif', np.asarray(instance_map_z).astype(np.uint16))
-            # imsave(save_dir + '/temp_seed.tif', np.asarray(seed_z).astype(np.float32))
-            instance_map = perform_ilp(instance_map_z, mean_object_size=mean_object_size)
+            instance_map = perform_ilp(
+                instance_map_z, mean_object_size=mean_object_size
+            )
 
-            if ('instance' in sample):
-                instances = sample['instance'].squeeze()  # Z Y X  (squeeze takes away first two dimensions) or DYX
+            if "instance" in sample:
+                instances = sample[
+                    "instance"
+                ].squeeze()  # Z Y X  (squeeze takes away first two dimensions) or DYX
 
-            base, _ = os.path.splitext(os.path.basename(sample['im_name'][0]))
+            base, _ = os.path.splitext(os.path.basename(sample["im_name"][0]))
             image_file_names.append(base)
 
-            if ('instance' in sample):
-                all_results = matching_dataset(y_true=[instances.cpu().detach().numpy()],
-                                               y_pred=[instance_map], thresh=ap_val,
-                                               show_progress=False)
-                if (verbose):
+            if "instance" in sample:
+                all_results = matching_dataset(
+                    y_true=[instances.cpu().detach().numpy()],
+                    y_pred=[instance_map],
+                    thresh=ap_val,
+                    show_progress=False,
+                )
+                if verbose:
                     print("Accuracy: {:.03f}".format(all_results.accuracy), flush=True)
                 result_list.append(all_results.accuracy)
 
             if save_images and ap_val == 0.5:
-
-                if not os.path.exists(os.path.join(save_dir, 'predictions/')):
-                    os.makedirs(os.path.join(save_dir, 'predictions/'))
-                    print("Created new directory {}".format(os.path.join(save_dir, 'predictions/')))
-                if not os.path.exists(os.path.join(save_dir, 'seeds/')):
-                    os.makedirs(os.path.join(save_dir, 'seeds/'))
-                    print("Created new directory {}".format(os.path.join(save_dir, 'seeds/')))
+                if not os.path.exists(os.path.join(save_dir, "predictions/")):
+                    os.makedirs(os.path.join(save_dir, "predictions/"))
+                    print(
+                        "Created new directory {}".format(
+                            os.path.join(save_dir, "predictions/")
+                        )
+                    )
+                if not os.path.exists(os.path.join(save_dir, "seeds/")):
+                    os.makedirs(os.path.join(save_dir, "seeds/"))
+                    print(
+                        "Created new directory {}".format(
+                            os.path.join(save_dir, "seeds/")
+                        )
+                    )
 
                 # save predictions
-                base, _ = os.path.splitext(os.path.basename(sample['im_name'][0]))
-                instances_file = os.path.join(save_dir, 'predictions/', base + '.tif')
+                base, _ = os.path.splitext(os.path.basename(sample["im_name"][0]))
+                instances_file = os.path.join(save_dir, "predictions/", base + ".tif")
                 imsave(instances_file, instance_map)
 
-                seeds_file = os.path.join(save_dir, 'seeds/', base + '.tif')
-                imsave(seeds_file, np.asarray(seed_z).astype(np.float32))
-
-        if save_results and 'instance' in sample:
-            if not os.path.exists(os.path.join(save_dir, 'results/')):
-                os.makedirs(os.path.join(save_dir, 'results/'))
-                print("Created new directory {}".format(os.path.join(save_dir, 'results/')))
-            txt_file = os.path.join(save_dir,
-                                    'results/combined_AP-' + '{:.02f}'.format(ap_val) + '_tta-' + str(tta) + '.txt')
-            with open(txt_file, 'w') as f:
+                seeds_file = os.path.join(save_dir, "seeds/", base + ".tif")
+                imsave(seeds_file, np.asarray(seed_z).astype(float))
+
+        if save_results and "instance" in sample:
+            if not os.path.exists(os.path.join(save_dir, "results/")):
+                os.makedirs(os.path.join(save_dir, "results/"))
+                print(
+                    "Created new directory {}".format(
+                        os.path.join(save_dir, "results/")
+                    )
+                )
+            txt_file = os.path.join(
+                save_dir,
+                "results/combined_AP-"
+                + "{:.02f}".format(ap_val)
+                + "_tta-"
+                + str(tta)
+                + ".txt",
+            )
+            with open(txt_file, "w") as f:
                 f.writelines(
-                    "image_file_name, min_mask_sum, min_unclustered_sum, min_object_size, seed_thresh, intersection_threshold, accuracy \n")
+                    "image_file_name, min_mask_sum, min_unclustered_sum, \
+                            min_object_size, seed_thresh, \
+                            intersection_threshold, accuracy \n"
+                )
                 f.writelines("+++++++++++++++++++++++++++++++++\n")
                 for ind, im_name in enumerate(image_file_names):
-                    im_name_tif = im_name + '.tif'
+                    im_name_tif = im_name + ".tif"
                     score = result_list[ind]
                     f.writelines(
-                        "{} {:.02f} {:.02f} {:.02f} {:.02f} {:.02f} {:.05f} \n".format(im_name_tif, min_mask_sum,
-                                                                                       min_unclustered_sum,
-                                                                                       min_object_size, seed_thresh,
-                                                                                       ap_val, score))
+                        "{} {:.02f} {:.02f} {:.02f} {:.02f} \
+                                {:.02f} {:.05f} \n".format(
+                            im_name_tif,
+                            min_mask_sum,
+                            min_unclustered_sum,
+                            min_object_size,
+                            seed_thresh,
+                            ap_val,
+                            score,
+                        )
+                    )
                 f.writelines("+++++++++++++++++++++++++++++++++\n")
-                f.writelines("Average Precision (AP_dsb)  {:.02f} {:.05f}\n".format(ap_val, np.mean(result_list)))
+                f.writelines(
+                    "Average Precision (AP_dsb)  {:.02f} {:.05f}\n".format(
+                        ap_val, np.mean(result_list)
+                    )
+                )
 
         if len(result_list) != 0:
             print(
-                "Mean Average Precision (AP_dsb) at IOU threshold = {} at seediness threshold = {:.05f}, is equal to {:.05f}".format(
-                    ap_val, seed_thresh, np.mean(result_list)))
+                "Mean Average Precision (AP_dsb) at IOU threshold \
+                        = {} at seediness threshold \
+                        = {:.05f}, is equal to {:.05f}".format(
+                    ap_val, seed_thresh, np.mean(result_list)
+                )
+            )
             return -np.mean(result_list)
         else:
             return 0.0
diff --git a/EmbedSeg/utils/test_time_augmentation.py b/EmbedSeg/utils/test_time_augmentation.py
index 2154322..6b87207 100644
--- a/EmbedSeg/utils/test_time_augmentation.py
+++ b/EmbedSeg/utils/test_time_augmentation.py
@@ -56,9 +56,12 @@ def process_flips(im_numpy):
 
     """
     im_numpy_correct = im_numpy
-    im_numpy_correct[0, 1, ...] = -1 * im_numpy[0, 1, ...]  # because flipping is always along y-axis, so only the y-offset gets affected
+    im_numpy_correct[0, 1, ...] = (
+        -1 * im_numpy[0, 1, ...]
+    )  # because flipping is always along y-axis, so only the y-offset gets affected
     return im_numpy_correct
 
+
 def to_cuda_3d(im_numpy):
     """
     Converts 3D Numpy Image on CPU to PyTorch Tensor on GPU, with an extra dimension
@@ -93,7 +96,7 @@ def apply_tta_2d(im, model):
     PyTorch Tensor on GPU (15YX)
 
     """
-    im_numpy = im.cpu().detach().numpy() # BCYX
+    im_numpy = im.cpu().detach().numpy()  # BCYX
     im0 = im_numpy[0, ...]  # remove batch dimension, now CYX
     im1 = np.rot90(im0, 1, (1, 2))
     im2 = np.rot90(im0, 2, (1, 2))
@@ -103,7 +106,7 @@ def apply_tta_2d(im, model):
     im6 = np.flip(im2, 1)
     im7 = np.flip(im3, 1)
 
-    im0_cuda = to_cuda(im0) # BCYX
+    im0_cuda = to_cuda(im0)  # BCYX
     im1_cuda = to_cuda(np.ascontiguousarray(im1))
     im2_cuda = to_cuda(np.ascontiguousarray(im2))
     im3_cuda = to_cuda(np.ascontiguousarray(im3))
@@ -195,9 +198,19 @@ def apply_tta_2d(im, model):
     output7_numpy_correct[0, 3, ...] = output7_numpy_flipped[0, 2, ...]
     output7_numpy_correct[0, 4, ...] = output7_numpy_flipped[0, 4, ...]
 
-    output = np.concatenate((output0_numpy_correct, output1_numpy_correct, output2_numpy_correct, output3_numpy_correct,
-                             output4_numpy_correct, output5_numpy_correct, output6_numpy_correct,
-                             output7_numpy_correct), 0)
+    output = np.concatenate(
+        (
+            output0_numpy_correct,
+            output1_numpy_correct,
+            output2_numpy_correct,
+            output3_numpy_correct,
+            output4_numpy_correct,
+            output5_numpy_correct,
+            output6_numpy_correct,
+            output7_numpy_correct,
+        ),
+        0,
+    )
     output = np.mean(output, 0, keepdims=True)  # 1 5 Y X
     return torch.from_numpy(output).float().cuda()
 
@@ -390,7 +403,6 @@ def apply_tta_3d(im, model, index):
         output_numpy_correct[0, 5, ...] = output_numpy_flipped[0, 5, ...]
         output_numpy_correct[0, 6, ...] = output_numpy_flipped[0, 6, ...]
 
-
     elif index == 10:
         temp = np.rot90(im_transformed, 2, (2, 3))
         temp = np.flip(temp, 1)
@@ -532,4 +544,4 @@ def apply_tta_3d(im, model, index):
         output_numpy_correct[0, 5, ...] = output_numpy_flipped[0, 5, ...]
         output_numpy_correct[0, 6, ...] = output_numpy_flipped[0, 6, ...]
 
-    return output_numpy_correct
\ No newline at end of file
+    return output_numpy_correct