diff --git a/README.md b/README.md
index 3315710d..033d4096 100644
--- a/README.md
+++ b/README.md
@@ -64,12 +64,12 @@ isc clusters # view the status of the clusters
 
 (from https://github.com/pytorch/vision/tree/main/references/segmentation)
 
-- WIP [fcn_resnet101.isc](./tv-segmentation/fcn_resnet101.isc)
-- WIP [deeplabv3_mobilenet_v3_large.isc](./tv-segmentation/deeplabv3_mobilenet_v3_large.isc)
+- [fcn_resnet101.isc](./tv-segmentation/fcn_resnet101.isc)
+- [deeplabv3_mobilenet_v3_large.isc](./tv-segmentation/deeplabv3_mobilenet_v3_large.isc)
 
 ### tv-detection
 
 (from https://github.com/pytorch/vision/tree/main/references/detection)
 
-- WIP [maskrcnn_resnet50_fpn.isc](./tv-detection/fasterrcnn_resnet50_fpn.isc)
-- WIP [retinanet_resnet50_fpn.isc](./tv-detection/retinanet_resnet50_fpn.isc)
+- [maskrcnn_resnet50_fpn.isc](./tv-detection/fasterrcnn_resnet50_fpn.isc)
+- [retinanet_resnet50_fpn.isc](./tv-detection/retinanet_resnet50_fpn.isc)
diff --git a/cycling_utils/cycling_utils/__init__.py b/cycling_utils/cycling_utils/__init__.py
index ddd8f60c..81a278f9 100644
--- a/cycling_utils/cycling_utils/__init__.py
+++ b/cycling_utils/cycling_utils/__init__.py
@@ -1,5 +1,5 @@
-from .saving import atomic_torch_save
-from .sampler import InterruptableDistributedSampler
-from .lightning_utils import EpochHandler
+from .timer import Timer, TimestampedTimer
+from .saving import atomic_torch_save, MetricsTracker
+from .sampler import InterruptableDistributedSampler, InterruptableDistributedGroupedBatchSampler
 
-__all__ = ["InterruptableDistributedSampler", "atomic_torch_save", "EpochHandler"]
\ No newline at end of file
+__all__ = ["InterruptableDistributedSampler", "InterruptableDistributedGroupedBatchSampler", "atomic_torch_save", "Timer", "TimestampedTimer"]
diff --git a/cycling_utils/cycling_utils/sampler.py b/cycling_utils/cycling_utils/sampler.py
index 097ad9c0..100ea811 100644
--- a/cycling_utils/cycling_utils/sampler.py
+++ b/cycling_utils/cycling_utils/sampler.py
@@ -2,6 +2,8 @@
 import torch
 from torch.utils.data import Dataset, DistributedSampler
 from contextlib import contextmanager
+from collections import defaultdict
+from itertools import chain, repeat
 
 class HasNotResetProgressError(Exception):
     pass
@@ -113,3 +115,167 @@ def in_epoch(self, epoch):
         self.set_epoch(epoch)
         yield
         self._reset_progress()
+
+def _repeat_to_at_least(iterable, n):
+    repeat_times = math.ceil(n / len(iterable))
+    repeated = chain.from_iterable(repeat(iterable, repeat_times))
+    return list(repeated)
+        
+class InterruptableDistributedGroupedBatchSampler(DistributedSampler):
+    def __init__(
+        self,
+        dataset: Dataset,
+        group_ids: list, 
+        batch_size: int,
+        num_replicas: int | None = None,
+        rank: int | None = None,
+        shuffle: bool = True,
+        seed: int = 0,
+        drop_last: bool = False,
+    ) -> None:
+        """
+        This is a DistributedSampler that can be suspended and resumed.
+
+        This works by keeping track of the sample batches that have already been 
+        dispatched. This InterruptableDistributedGroupedBatchSampler also 
+        reproduces the sampling strategy exhibited in the torch vision detection
+        reference wherein batches are created from images from within the same
+        'group', defined in the torchvision example by similarity of image 
+        aspect ratio. 
+
+        https://github.com/pytorch/vision/tree/main/references/detection
+
+        For this reason, InterruptableDistributedGroupedBatchSampler progress is
+        tracked in units of batches, not samples. This is an important
+        distinction from the InterruptableDistributedSampler which tracks progress
+        in units of samples. The progress is reset to 0 at the end of each epoch.
+
+        The epoch is set to 0 at initialization and incremented at the start 
+        of each epoch.
+
+        Suspending and resuming the sampler is done by saving and loading the
+        state dict. The state dict contains the epoch and progress.
+        """
+        super().__init__(dataset, num_replicas, rank, shuffle, seed, drop_last)
+
+        # OVERALL STATUS INDICATOR
+        self.progress = 0
+        self._has_reset_progress = True
+        self.batch_size = batch_size
+        self.group_ids = group_ids
+        self.batches = self._create_batches()
+
+    def _create_batches(self):
+        if self.shuffle:
+            # deterministically shuffle based on epoch and seed
+            g = torch.Generator()
+            g.manual_seed(self.seed + self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]
+        else:
+            indices = list(range(len(self.dataset)))  # type: ignore[arg-type]
+
+        if not self.drop_last:
+            # add extra samples to make dataset evenly divisible accross ranks
+            padding_size = self.total_size - len(indices)
+            if padding_size <= len(indices):
+                indices += indices[:padding_size]
+            else:
+                indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]
+        else:
+            # remove tail of data to make dataset evenly divisible accross ranks
+            indices = indices[: self.total_size]
+        assert len(indices) == self.total_size
+
+        # subsample indices to use on this rank
+        indices = indices[self.rank : self.total_size : self.num_replicas]
+        assert len(indices) == self.num_samples
+
+        # PRE-COMPUTE GROUPED BATCHES
+        buffer_per_group = defaultdict(list)
+        samples_per_group = defaultdict(list)
+        self.num_batches = math.ceil(len(indices)/ self.batch_size)
+
+        batches = [] # pre-computed so progress refers to batches, not samples.
+        for idx in indices:
+            group_id = self.group_ids[idx]
+            buffer_per_group[group_id].append(idx)
+            samples_per_group[group_id].append(idx)
+            if len(buffer_per_group[group_id]) == self.batch_size:
+                batches.append(buffer_per_group[group_id])
+                del buffer_per_group[group_id]
+            assert len(buffer_per_group[group_id]) < self.batch_size
+
+        # now we have run out of elements that satisfy
+        # the group criteria, let's return the remaining
+        # elements so that the size of the sampler is
+        # deterministic
+        num_remaining = self.num_batches - len(batches)
+        if num_remaining > 0:
+            # for the remaining batches, take first the buffers with the largest number
+            # of elements
+            for group_id, _ in sorted(buffer_per_group.items(), key=lambda x: len(x[1]), reverse=True):
+                remaining = self.batch_size - len(buffer_per_group[group_id])
+                samples_from_group_id = _repeat_to_at_least(samples_per_group[group_id], remaining)
+                buffer_per_group[group_id].extend(samples_from_group_id[:remaining])
+                assert len(buffer_per_group[group_id]) == self.batch_size
+                batches.append(buffer_per_group[group_id])
+                num_remaining -= 1
+                if num_remaining == 0:
+                    break
+
+        # Check that the batches are all good to go
+        assert len(batches) == self.num_batches
+        return batches
+
+    def _reset_progress(self):
+        self.progress = 0
+        self._has_reset_progress = True
+
+    def set_epoch(self, epoch: int) -> None:
+        raise NotImplementedError("Use `with sampler.in_epoch(epoch)` instead of `sampler.set_epoch(epoch)`")
+
+    def _set_epoch(self, epoch):
+        if not self._has_reset_progress:
+            raise HasNotResetProgressError("You must reset progress before setting epoch e.g. `sampler.reset_progress()`\nor use `with sampler.in_epoch(epoch)` instead of `sampler.set_epoch(epoch)`")
+        self.epoch = epoch
+
+    def state_dict(self):
+        return {"progress": self.progress, "epoch": self.epoch}
+
+    def load_state_dict(self, state_dict):
+        self.progress = state_dict["progress"]
+        if not self.progress <= self.num_batches:
+            raise AdvancedTooFarError(f"progress should be less than or equal to the number of batches. progress: {self.progress}, num_batches: {self.num_batches}")
+        self.epoch = state_dict["epoch"]
+
+    def advance(self):
+        """
+        Record that one batch has been consumed.
+        """
+        self.progress += 1
+        if self.progress > self.num_batches:
+            raise AdvancedTooFarError(f"You have advanced too far. You can only advance up to the total number of batches: {self.num_batches}.")
+
+    def __iter__(self):
+
+        # slice from progress to pick up where we left off
+        for batch in self.batches[self.progress:]:
+            yield batch
+
+    def __len__(self):
+        return self.num_batches
+
+    @contextmanager
+    def in_epoch(self, epoch):
+        """
+        This context manager is used to set the epoch. It is used like this:
+        ```
+        for epoch in range(0, 10):
+            with sampler.in_epoch(epoch):
+                for step, (x, ) in enumerate(dataloader):
+                    # work would be done here...
+        ```
+        """
+        self._set_epoch(epoch)
+        yield
+        self._reset_progress()
diff --git a/cycling_utils/cycling_utils/saving.py b/cycling_utils/cycling_utils/saving.py
index a71fef99..17f91898 100644
--- a/cycling_utils/cycling_utils/saving.py
+++ b/cycling_utils/cycling_utils/saving.py
@@ -1,9 +1,87 @@
 from pathlib import Path
 import os
 import torch
+import torch.distributed as dist
+from collections import defaultdict
 
-def atomic_torch_save(obj, f: str | Path, **kwargs):
+def atomic_torch_save(obj, f: str | Path, timer=None, **kwargs):
     f = str(f)
     temp_f = f + ".temp"
     torch.save(obj, temp_f, **kwargs)
-    os.replace(temp_f, f)
\ No newline at end of file
+    if timer is not None:
+        timer.report(f'saving temp checkpoint')
+    os.replace(temp_f, f)
+    if timer is not None:
+        timer.report(f'replacing temp checkpoint with checkpoint')
+        return timer
+    else:
+        return
+
+class MetricsTracker:
+    '''
+    This is a general purpose MetricsTracker to assist with recording metrics from
+    a disributed cluster.
+
+    The MetricsTracker is initialised without any prior knowledge of the metrics
+    to be tracked.
+
+    >>> metrics = MetricsTracker()
+
+    Metrics can be accumulated as required, for example after each batch is procesed
+    by the model, by passing a dictionary with metrics to be updated, then reduced 
+    accross all nodes. Metric values are stored in a defaultdict.
+
+    >>> preds = model(input)
+    >>> loss = loss_fn(preds, targs)
+    >>> metrics.update({"images_seen": len(images), "loss": loss.item()})
+    >>> metrics.reduce()
+
+    Metrics are assumed to be summable scalar values. After calling reduce(), the 
+    metrics.local object contains the sum of corresponding metrics from all nodes
+    which can be used for intermediate reporting or logging.
+
+    >>> writer = SummaryWriter()
+    >>> for metric,val in metrics.local.items():
+    >>>     writer.add_scalar(metric, val, step)
+    >>> writer.flush()
+    >>> writer.close()
+
+    Once all processing of the current batch has been completed, the MetricsTracker
+    can be prepared for the next batch using reset_local().
+
+    >>> metrics.reset_loca()
+
+    Metrics are also accumulated for consecutive batches in the metrics.agg object.
+    At the end of an epoch the MetricsTracker can be reset using end_epoch().
+
+    >>> metrics.end_epoch()
+
+    The MetricsTracker saves a copy of the accumulated metrics (metrics.agg) for
+    each epoch which can be stored within a checkpoint.
+    '''
+    def __init__(self):
+        self.local = defaultdict(float)
+        self.agg = defaultdict(float)
+        self.epoch_reports = []
+
+    def update(self, metrics: dict):
+        for m,v in metrics.items():
+            self.local[m] += v
+        
+    def reduce(self):
+        names, local = zip(*self.local.items())
+        local = torch.tensor(local, dtype=torch.float16, requires_grad=False, device='cuda')
+        dist.all_reduce(local, op=dist.ReduceOp.SUM)
+        self.local = defaultdict(float, zip(names, local.cpu().numpy()))
+        for k in self.local:
+            self.agg[k] += self.local[k]
+
+    def reset_local(self):
+        self.local = defaultdict(float)
+    
+    def end_epoch(self):
+        self.epoch_reports.append(dict(self.agg))
+        self.local = defaultdict(float)
+        self.agg = defaultdict(float)
+    
+
diff --git a/cycling_utils/cycling_utils/timer.py b/cycling_utils/cycling_utils/timer.py
new file mode 100644
index 00000000..6915b313
--- /dev/null
+++ b/cycling_utils/cycling_utils/timer.py
@@ -0,0 +1,64 @@
+import os, time
+from datetime import datetime
+
+class Timer:
+    '''
+    This Timer can be integrated within a training routine to provide point-to-point
+    script timing and reporting.
+
+    def main():
+        timer = Timer()
+        time.sleep(2)
+        timer.report("sleeping for 2 seconds")
+        time.sleep(3)
+        timer.report("sleeping for 3 seconds")
+
+    >>> main()
+    Start                                       0.000 ms     0.000 s total
+    Completed sleeping for 2 seconds        2,000.000 ms     2.000 s total
+    Completed sleeping for 3 seconds        3,000.000 ms     5.000 s total
+    '''
+    def __init__(self, report=None, start_time=None, running=0):
+        self.start_time = start_time if start_time is not None else time.time()
+        self.running = running
+        if str(os.environ["RANK"]) == "0":
+            report = report if report else "Start"
+            print("[{:<80}] {:>12} ms, {:>12} s total".format(report, f'{0.0:,.3f}', f'{0.0:,.2f}'))
+    def report(self, annot):
+        if str(os.environ["RANK"]) == "0":
+            now = time.time()
+            duration = now - self.start_time
+            self.running += duration
+            print("Completed {:<70}{:>12} ms, {:>12} s total".format(annot, f'{1000*duration:,.3f}', f'{self.running:,.2f}'))
+            self.start_time = now
+
+class TimestampedTimer:
+    '''
+    This TimestampedTimer can be integrated within a training routine to provide 
+    point-to-point script timing and reporting.
+
+    def main():
+        timer = TimestampedTimer()
+        time.sleep(2)
+        timer.report("sleeping for 2 seconds")
+        time.sleep(3)
+        timer.report("sleeping for 3 seconds")
+
+    >>> main()
+    [TIME] Start                                       0.000 ms     0.000 s total
+    [TIME] Completed sleeping for 2 seconds        2,000.000 ms     2.000 s total
+    [TIME] Completed sleeping for 3 seconds        3,000.000 ms     5.000 s total
+    '''
+    def __init__(self, report=None, start_time=None, running=0):
+        if str(os.environ.get("RANK","NONE")) in ["0", "NONE"]:
+            self.start_time = start_time if start_time is not None else time.time()
+            self.running = running
+            report = report if report else "Start"
+            print("[ {} ] Completed {:<70}{:>12} ms, {:>12} s total".format(time.strftime("%Y-%m-%d %H:%M:%S"), report, f'{0.0:,.3f}', f'{0.0:,.2f}'))
+    def report(self, annot):
+        if str(os.environ.get("RANK","NONE")) in ["0", "NONE"]:
+            now = time.time()
+            duration = now - self.start_time
+            self.running += duration
+            print("[ {} ] Completed {:<70}{:>12} ms, {:>12} s total".format(time.strftime("%Y-%m-%d %H:%M:%S"), annot, f'{1000*duration:,.3f}', f'{self.running:,.2f}'))
+            self.start_time = now
\ No newline at end of file
diff --git a/monai_brats_mri_2d/README.md b/monai_brats_mri_2d/README.md
new file mode 100644
index 00000000..d57caf6e
--- /dev/null
+++ b/monai_brats_mri_2d/README.md
@@ -0,0 +1,3 @@
+# MONAI Generative Models Installation
+For this demonstration, you will need to clone the MONAI GenerativeModels GitHub repository and follow the instructions for installation. This will install the `generative` package from MONAI.
+You will then need to run `pip install -r requirements-dev.txt` to install other necessary dependencies. You may then also need to ensure that monai version 1.2.0 is installed using the command `pip install monai==1.2.0` as later versions of monai do not support all of the transforms used in this example.
\ No newline at end of file
diff --git a/monai_brats_mri_2d/brats_mri_2d_diff.isc b/monai_brats_mri_2d/brats_mri_2d_diff.isc
new file mode 100644
index 00000000..a7ad9df0
--- /dev/null
+++ b/monai_brats_mri_2d/brats_mri_2d_diff.isc
@@ -0,0 +1,6 @@
+experiment_name="brats_mri_2d_diff"
+gpu_type="24GB VRAM GPU"
+nnodes = 11
+venv_path = "~/.venv/bin/activate"
+output_path = "~outputs/brats_mri_2d_diff"
+command="train_cycling_diff.py --data-path=/mnt/.node1/Open-Datsets/MONAI --resume $OUTPUT_PATH/checkpoint.isc, --gen-load-path ~/output_brats_mri_2d_gen/exp_1855/checkpoint.isc --tboard-path $OUTPUT_PATH/tb"
\ No newline at end of file
diff --git a/monai_brats_mri_2d/brats_mri_2d_gen.isc b/monai_brats_mri_2d/brats_mri_2d_gen.isc
new file mode 100644
index 00000000..78e26639
--- /dev/null
+++ b/monai_brats_mri_2d/brats_mri_2d_gen.isc
@@ -0,0 +1,6 @@
+experiment_name="brats_mri_2d_gen"
+gpu_type="24GB VRAM GPU"
+nnodes = 11
+venv_path = "~/.venv/bin/activate"
+output_path = "~/outputs/brats_mri_2d_gen"
+command="train_cycling_gen.py --lr 1e-5 --data-path=/mnt/.node1/Open-Datasets/MONAI --resume $OUTPUT_PATH/checkpoint.isc --tboard-path $OUTPUT_PATH/tb --prev-resume /mnt/Client/StrongHumans/strong_adam/outputs/brats_mri_2d_gen/301e7ac7-0c9a-4daa-920e-57ea5ea983b9/checkpoint.isc"
\ No newline at end of file
diff --git a/monai_brats_mri_2d/loops.py b/monai_brats_mri_2d/loops.py
new file mode 100644
index 00000000..ad6fa2a4
--- /dev/null
+++ b/monai_brats_mri_2d/loops.py
@@ -0,0 +1,388 @@
+import torch, utils
+from torch.cuda.amp import autocast
+import torch.nn.functional as F
+from cycling_utils import atomic_torch_save
+from generative.losses.adversarial_loss import PatchAdversarialLoss
+from torch.utils.tensorboard import SummaryWriter
+from torchvision.utils import make_grid
+
+## -- AUTO-ENCODER - ##
+
+def compute_kl_loss(z_mu, z_sigma):
+    kl_loss = 0.5 * torch.sum(
+        z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, 
+        dim=list(range(1, len(z_sigma.shape)))
+    )
+    return torch.sum(kl_loss) / kl_loss.shape[0]
+
+intensity_loss = torch.nn.L1Loss()
+adv_loss = PatchAdversarialLoss(criterion="least_squares")
+
+def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, perceptual_loss, kl_weight, perceptual_weight, adv_weight):
+    # Image intrinsic qualities
+    recons_loss = intensity_loss(gen_images, real_images)
+    kl_loss = compute_kl_loss(z_mu, z_sigma)
+    p_loss = perceptual_loss(gen_images.float(), real_images.float())
+    loss_g = recons_loss + (kl_weight * kl_loss) + (perceptual_weight * p_loss)
+    # Discrimnator-based loss
+    logits_fake = disc_net(gen_images)[-1]
+    generator_loss = adv_loss(logits_fake, target_is_real=True, for_discriminator=False)
+    loss_g = loss_g + (adv_weight * generator_loss)
+    return loss_g
+
+def discriminator_loss(gen_images, real_images, disc_net, adv_weight):
+    logits_fake = disc_net(gen_images.contiguous().detach())[-1]
+    loss_d_fake = adv_loss(logits_fake, target_is_real=False, for_discriminator=True)
+    logits_real = disc_net(real_images.contiguous().detach())[-1]
+    loss_d_real = adv_loss(logits_real, target_is_real=True, for_discriminator=True)
+    discriminator_loss = (loss_d_fake + loss_d_real) * 0.5
+    loss_d = adv_weight * discriminator_loss
+    return loss_d
+
+def train_generator_one_epoch(
+        args, epoch, generator, discriminator, optimizer_g, optimizer_d, train_sampler, val_sampler,
+        scaler_g, scaler_d, train_loader, perceptual_loss, device, timer, metrics
+    ):
+
+    generator.train()
+    discriminator.train()
+
+    train_step = train_sampler.progress // train_loader.batch_size
+    total_steps = int(len(train_sampler) / train_loader.batch_size)
+    print(f'\nTraining / resuming epoch {epoch} from training step {train_step}\n')
+
+    for batch in train_loader:
+
+        images = batch["image"].to(device)
+        timer.report(f'train batch {train_step} to device')
+
+        # TRAIN GENERATOR
+
+        optimizer_g.zero_grad(set_to_none=True)
+
+        with autocast(enabled=True):
+
+            reconstruction, z_mu, z_sigma = generator(images)
+            timer.report(f'train batch {train_step} generator forward')
+
+            loss_g = generator_loss(
+                reconstruction, images, z_mu, z_sigma, discriminator, perceptual_loss, 
+                args.kl_weight, args.perceptual_weight, args.adv_weight
+            )
+            timer.report(f'train batch {train_step} generator loss: {loss_g.item():.3f}')
+
+        scaler_g.scale(loss_g).backward()
+        scaler_g.step(optimizer_g)
+        scaler_g.update()
+        timer.report(f'train batch {train_step} generator backward')
+
+        # TRAIN DISCRIMINATOR
+
+        optimizer_d.zero_grad(set_to_none=True)
+
+        with autocast(enabled=True):
+
+            loss_d = discriminator_loss(
+                reconstruction, images, discriminator, args.adv_weight
+            )
+            timer.report(f'train batch {train_step} discriminator loss {loss_d.item():.3f}')
+
+        scaler_d.scale(loss_d).backward()
+        scaler_d.step(optimizer_d)
+        scaler_d.update()
+        timer.report(f'train batch {train_step} discriminator backward')
+
+        # Reduce metrics accross nodes
+        metrics["train"].update({"train_images_seen":len(images), "loss_g":loss_g.item(), "loss_d": loss_d.item()})
+        metrics["train"].reduce()
+
+        gen_loss = metrics["train"].local["loss_g"] / metrics["train"].local["train_images_seen"]
+        disc_loss = metrics["train"].local["loss_d"] / metrics["train"].local["train_images_seen"]
+        print("Epoch [{}] Step [{}/{}], gen_loss: {:.3f}, disc_loss: {:.3f}".format(epoch, train_step, total_steps, gen_loss, disc_loss))
+
+        metrics["train"].reset_local()
+
+        timer.report(f'train batch {train_step} metrics update')
+
+        ## Checkpointing
+        print(f"Saving checkpoint at epoch {epoch} train batch {train_step}")
+        train_sampler.advance(len(images))
+        train_step = train_sampler.progress // train_loader.batch_size
+        
+        if train_step == total_steps:
+            metrics["train"].end_epoch()
+
+        if utils.is_main_process() and train_step % 1 == 0: # Checkpointing every batch
+            writer = SummaryWriter(log_dir=args.tboard_path)
+            writer.add_scalar("Train/gen_loss", gen_loss, train_step + epoch * total_steps)
+            writer.add_scalar("Train/disc_loss", disc_loss, train_step + epoch * total_steps)
+            writer.flush()
+            writer.close()
+            checkpoint = {
+                # Universals
+                "args": args,
+                "epoch": epoch,
+                # State variables
+                "generator": generator.module.state_dict(),
+                "discriminator": discriminator.module.state_dict(),
+                "optimizer_g": optimizer_g.state_dict(),
+                "optimizer_d": optimizer_d.state_dict(),
+                "scaler_g": scaler_g.state_dict(),
+                "scaler_d": scaler_d.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "val_sampler": val_sampler.state_dict(),
+                # Metrics
+                "metrics": metrics,
+            }
+            timer = atomic_torch_save(checkpoint, args.resume, timer)
+
+    gen_loss = metrics["train"].epoch_reports[-1]["loss_g"] / metrics["train"].epoch_reports[-1]["train_images_seen"]
+    disc_loss = metrics["train"].epoch_reports[-1]["loss_d"] / metrics["train"].epoch_reports[-1]["train_images_seen"]
+    print("Epoch [{}] :: gen_loss: {:,.3f}, disc_loss: {:,.3f}".format(epoch, gen_loss, disc_loss))
+    return generator, timer, metrics
+
+
+def evaluate_generator(
+        args, epoch, generator, discriminator, optimizer_g, optimizer_d, train_sampler, val_sampler,
+        scaler_g, scaler_d, val_loader, device, timer, metrics
+    ):
+
+    generator.eval()
+
+    val_step = val_sampler.progress // val_loader.batch_size
+    total_steps = int(len(val_sampler) / val_loader.batch_size)
+    print(f'\nEvaluating / resuming epoch {epoch} from eval step {val_step}\n')
+
+    with torch.no_grad():
+        for batch in val_loader:
+
+            images = batch["image"].to(device)
+            timer.report(f'eval batch {val_step} to device')
+
+            with autocast(enabled=True):
+
+                reconstruction, _, _ = generator(images)
+                timer.report(f'eval batch {val_step} forward')
+                recons_loss = F.l1_loss(images.float(), reconstruction.float())
+                timer.report(f'eval batch {val_step} recons_loss')
+
+            metrics["val"].update({"val_images_seen": len(images), "val_loss": recons_loss.item()})
+            metrics["val"].reduce()
+            metrics["val"].reset_local()
+
+            timer.report(f'eval batch {val_step} metrics update')
+
+            ## Checkpointing
+            print(f"Saving checkpoint at epoch {epoch} val batch {val_step}")
+            val_sampler.advance(len(images))
+            val_step = val_sampler.progress // val_loader.batch_size
+
+            if val_step == total_steps:
+
+                val_loss = metrics["val"].agg["val_loss"] / metrics["val"].agg["val_images_seen"]
+                if utils.is_main_process():
+
+                    writer = SummaryWriter(log_dir=args.tboard_path)
+                    writer.add_scalar("Val/loss", val_loss, epoch)
+
+                    # images_list = torch.zeros((11*6, *images.shape[1:]), device=device, dtype=images.dtype)
+                    # reconstruction_list = torch.zeros((11*6, *reconstruction.shape[1:]), device=device, dtype=reconstruction.dtype)
+                    # dist.all_gather_into_tensor(images_list, images.clone())
+                    # dist.all_gather_into_tensor(reconstruction_list, reconstruction)
+                    # plottable = torch.cat((images_list[0:5],reconstruction_list[0:5]))
+                    # plottable = (plottable * 255).to(torch.uint8)
+                    plottable = torch.cat((images, reconstruction))
+                    grid = make_grid(plottable, nrow=2)
+                    writer.add_image('Val/images', grid, epoch)
+
+                    writer.flush()
+                    writer.close()
+
+                print(f"Epoch {epoch} val loss: {val_loss:.4f}")
+                metrics["val"].end_epoch()
+
+            if utils.is_main_process() and val_step % 1 == 0: # Checkpointing every batch
+                checkpoint = {
+                    # Universals
+                    "args": args,
+                    "epoch": epoch,
+                    # State variables
+                    "generator": generator.module.state_dict(),
+                    "discriminator": discriminator.module.state_dict(),
+                    "optimizer_g": optimizer_g.state_dict(),
+                    "optimizer_d": optimizer_d.state_dict(),
+                    "scaler_g": scaler_g.state_dict(),
+                    "scaler_d": scaler_d.state_dict(),
+                    "train_sampler": train_sampler.state_dict(),
+                    "val_sampler": val_sampler.state_dict(),
+                    # Metrics
+                    "metrics": metrics,
+                }
+                timer = atomic_torch_save(checkpoint, args.resume, timer)
+
+    return timer, metrics
+
+
+## -- DIFFUSION MODEL - ##
+
+def train_diffusion_one_epoch(
+        args, epoch, unet, generator, optimizer_u, scaler_u, inferer, train_loader, val_loader, 
+        train_sampler, val_sampler, lr_scheduler, device, timer, metrics
+    ):
+
+    unet.train()
+    generator.eval()
+
+    train_step = train_sampler.progress // train_loader.batch_size
+    total_steps = len(train_sampler) // train_loader.batch_size
+    print(f'\nTraining / resuming epoch {epoch} from training step {train_step}\n')
+
+    for step, batch in enumerate(train_loader):
+
+        images = batch["image"].to(device)
+        timer.report(f'train batch {train_step} to device')
+
+        optimizer_u.zero_grad(set_to_none=True)
+
+        with autocast(enabled=True):
+
+            z_mu, z_sigma = generator.encode(images)
+            timer.report(f'train batch {train_step} generator encoded')
+            z = generator.sampling(z_mu, z_sigma)
+            timer.report(f'train batch {train_step} generator sampling')
+            noise = torch.randn_like(z).to(device)
+            timer.report(f'train batch {train_step} noise')
+            timesteps = torch.randint(0, inferer.scheduler.num_train_timesteps, (z.shape[0],), device=z.device).long()
+            timer.report(f'train batch {train_step} timesteps')
+            noise_pred = inferer(inputs=images, diffusion_model=unet, noise=noise, timesteps=timesteps, autoencoder_model=generator)
+            timer.report(f'train batch {train_step} noise_pred')
+            loss = F.mse_loss(noise_pred.float(), noise.float())
+            timer.report(f'train batch {train_step} loss')
+
+        scaler_u.scale(loss).backward()
+        scaler_u.step(optimizer_u)
+        scaler_u.update()
+        lr_scheduler.step()
+        timer.report(f'train batch {train_step} unet backward')
+
+        # Reduce metrics accross nodes
+        metrics["train"].update({"images_seen":len(images), "loss":loss.item()})
+        metrics["train"].reduce()
+
+        recons_loss = metrics["train"].local["loss"] / metrics["train"].local["images_seen"]
+        print("Epoch [{}] Step [{}/{}] :: loss: {:,.3f}".format(epoch, train_step, total_steps, recons_loss))
+
+        metrics["train"].reset_local()
+
+        timer.report(f'train batch {train_step} metrics update')
+
+        ## Checkpointing
+        print(f"Saving checkpoint at epoch {epoch} train batch {train_step}")
+        train_sampler.advance(len(images))
+        train_step = train_sampler.progress // train_loader.batch_size
+
+        if train_step == total_steps:
+            metrics["train"].end_epoch()
+
+        if utils.is_main_process() and train_step % 1 == 0: # Checkpointing every batch
+
+            writer = SummaryWriter(log_dir=args.tboard_path)
+            writer.add_scalar("Train/loss", recons_loss, train_step + epoch * total_steps)
+            writer.flush()
+            writer.close()
+
+            checkpoint = {
+                # Universals
+                "args": args,
+                "epoch": epoch,
+                # State variables
+                "unet": unet.module.state_dict(),
+                "optimizer_u": optimizer_u.state_dict(),
+                "scaler_u": scaler_u.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "val_sampler": val_sampler.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                # Metrics
+                "metrics": metrics,
+            }
+            timer = atomic_torch_save(checkpoint, args.resume, timer)
+
+    train_loss = metrics["train"].epoch_reports[-1]["loss"] / metrics["train"].epoch_reports[-1]["images_seen"]
+    print("Epoch [{}] :: epoch_loss: {:,.3f}".format(epoch, train_loss))
+    return unet, timer, metrics
+
+
+def evaluate_diffusion(
+        args, epoch, unet, generator, optimizer_u, scaler_u, inferer, train_loader, val_loader, 
+        train_sampler, val_sampler, lr_scheduler, device, timer, metrics
+    ):
+
+        unet.eval()
+
+        val_step = val_sampler.progress // val_loader.batch_size
+        total_steps = int(len(val_sampler) / val_loader.batch_size)
+        print(f'\nEvaluating / resuming epoch {epoch} from eval step {val_step}\n')
+
+        with torch.no_grad():
+            for step, batch in enumerate(val_loader):
+
+                images = batch["image"].to(device)
+                timer.report(f'eval batch {val_step} to device')
+
+                with autocast(enabled=True):
+
+                    z_mu, z_sigma = generator.encode(images)
+                    timer.report(f'eval batch {val_step} generator encoded')
+                    z = generator.sampling(z_mu, z_sigma)
+                    timer.report(f'eval batch {val_step} generator sampling')
+                    noise = torch.randn_like(z).to(device)
+                    timer.report(f'eval batch {val_step} noise')
+                    timesteps = torch.randint(0, inferer.scheduler.num_train_timesteps, (z.shape[0],), device=z.device).long()
+                    timer.report(f'eval batch {val_step} timesteps')
+                    noise_pred = inferer(inputs=images,diffusion_model=unet,noise=noise,timesteps=timesteps,autoencoder_model=generator)
+                    timer.report(f'eval batch {val_step} noise_pred')
+                    loss = F.mse_loss(noise_pred.float(), noise.float())
+                    timer.report(f'eval batch {val_step} loss')
+
+                metrics["val"].update({"images_seen": len(images), "loss": loss.item()})
+                metrics["val"].reduce()
+                metrics["val"].reset_local()
+
+                timer.report(f'eval batch {val_step} metrics update')
+
+                ## Checkpointing
+                print(f"Saving checkpoint at epoch {epoch} val batch {val_step}")
+                val_sampler.advance(len(images))
+                val_step = val_sampler.progress // val_loader.batch_size
+
+                if val_step == total_steps:
+                    metrics["val"].end_epoch()
+
+                if utils.is_main_process() and val_step % 1 == 0: # Checkpointing every batch
+                    print(f"Saving checkpoint at epoch {epoch} train batch {val_step}")
+                    checkpoint = {
+                        # Universals
+                        "args": args,
+                        "epoch": epoch,
+                        # State variables
+                        "unet": unet.module.state_dict(),
+                        "optimizer_u": optimizer_u.state_dict(),
+                        "scaler_u": scaler_u.state_dict(),
+                        "train_sampler": train_sampler.state_dict(),
+                        "val_sampler": val_sampler.state_dict(),
+                        "lr_scheduler": lr_scheduler.state_dict(),
+                        # Metrics
+                        "metrics": metrics,
+                    }
+                    timer = atomic_torch_save(checkpoint, args.resume, timer)
+
+        # val_loss = metrics["val"].agg[metrics["val"].map["val_loss"]] / metrics["val"].agg[metrics["val"].map["val_images_seen"]]
+        val_loss = metrics["val"].epoch_reports[-1]["loss"] / metrics["val"].epoch_reports[-1]["images_seen"]
+        if utils.is_main_process():
+            writer = SummaryWriter(log_dir=args.tboard_path)
+            writer.add_scalar("Val/loss", val_loss, epoch)
+            writer.flush()
+            writer.close()
+        print(f"Epoch [{epoch}] :: diff val loss: {val_loss:.4f}")
+
+        return timer, metrics
diff --git a/monai_brats_mri_2d/prep.py b/monai_brats_mri_2d/prep.py
new file mode 100644
index 00000000..a0e050e2
--- /dev/null
+++ b/monai_brats_mri_2d/prep.py
@@ -0,0 +1,23 @@
+# Ensuring required monai version is installed
+# pip install monai==1.2.0
+
+# Download the toy dataset from MONAI
+print("Downloadning BraTS2016/17")
+from monai.apps import DecathlonDataset
+from generative.losses.perceptual import PerceptualLoss
+
+# _ = DecathlonDataset(root_dir="/mnt/Datasets/Open-Datasets/MONAI", task="Task01_BrainTumour", section="training", download=True)
+# _ = DecathlonDataset(root_dir="/mnt/.node1/Open-Datsets/MONAI", task="Task01_BrainTumour", section="training", download=True)
+
+perceptual_loss = PerceptualLoss(
+    spatial_dims=2, network_type="resnet50", pretrained=True, #ImageNet pretrained weights used
+)
+
+# # Download the bigger dataset from Synapse
+# print("Downloadning BraTS2023")
+# import synapseclient
+# syn = synapseclient.Synapse()
+# syn.login('adam_peaston','AXXXXXXXXX2')
+# syn51514132 = syn.get(entity='syn51514132', downloadFile=True, downloadLocation="/mnt/Datasets/strongcompute_adam/MONAI", ifcollision="overwrite.local")
+# filepath = syn51514132.path
+# print(f"BraTS2023-GLI downloaded to {filepath}")
\ No newline at end of file
diff --git a/monai_brats_mri_2d/requirements-dev.txt b/monai_brats_mri_2d/requirements-dev.txt
new file mode 100644
index 00000000..c772949e
--- /dev/null
+++ b/monai_brats_mri_2d/requirements-dev.txt
@@ -0,0 +1,57 @@
+# Full requirements for developments
+-r requirements-min.txt
+pytorch-ignite==0.4.10
+gdown>=4.4.0
+scipy
+itk>=5.2
+nibabel
+pillow!=8.3.0  # https://github.com/python-pillow/Pillow/issues/5571
+tensorboard>=2.6  # https://github.com/Project-MONAI/MONAI/issues/5776
+scikit-image>=0.19.0
+tqdm>=4.47.0
+lmdb
+flake8>=3.8.1
+flake8-bugbear
+flake8-comprehensions
+flake8-executable
+pylint!=2.13  # https://github.com/PyCQA/pylint/issues/5969
+mccabe
+pep8-naming
+pycodestyle
+pyflakes
+black
+isort
+pytype>=2020.6.1; platform_system != "Windows"
+types-pkg_resources
+mypy>=0.790
+ninja
+# torchvision
+psutil
+Sphinx==3.5.3
+recommonmark==0.6.0
+sphinx-autodoc-typehints==1.11.1
+sphinx-rtd-theme==0.5.2
+cucim==22.8.1; platform_system == "Linux"
+openslide-python==1.1.2
+imagecodecs; platform_system == "Linux" or platform_system == "Darwin"
+tifffile; platform_system == "Linux" or platform_system == "Darwin"
+pandas
+requests
+einops
+transformers<4.22  # https://github.com/Project-MONAI/MONAI/issues/5157
+mlflow
+matplotlib!=3.5.0
+tensorboardX
+types-PyYAML
+pyyaml
+fire
+jsonschema
+pynrrd
+pre-commit
+pydicom
+h5py
+nni
+optuna
+git+https://github.com/Project-MONAI/MetricsReloaded@monai-support#egg=MetricsReloaded
+lpips==0.1.4
+xformers==0.0.16
diff --git a/monai_brats_mri_2d/requirements-min.txt b/monai_brats_mri_2d/requirements-min.txt
new file mode 100644
index 00000000..ceb9e346
--- /dev/null
+++ b/monai_brats_mri_2d/requirements-min.txt
@@ -0,0 +1,5 @@
+# Requirements for minimal tests
+-r requirements.txt
+setuptools>65.5.0,<66.0.0
+coverage>=5.5
+parameterized
diff --git a/monai_brats_mri_2d/requirements.txt b/monai_brats_mri_2d/requirements.txt
new file mode 100644
index 00000000..3f1ff86c
--- /dev/null
+++ b/monai_brats_mri_2d/requirements.txt
@@ -0,0 +1,3 @@
+numpy>=1.17
+torch>=1.8
+monai>=1.2.0rc1
diff --git a/monai_brats_mri_2d/train_cycling_diff.py b/monai_brats_mri_2d/train_cycling_diff.py
new file mode 100644
index 00000000..823ea3a9
--- /dev/null
+++ b/monai_brats_mri_2d/train_cycling_diff.py
@@ -0,0 +1,216 @@
+from cycling_utils import TimestampedTimer
+
+timer = TimestampedTimer()
+timer.report('importing Timer')
+
+import os
+import torch
+# import torch.nn.functional as F
+!pip install monai==1.2.0
+from monai import transforms
+from monai.apps import DecathlonDataset
+# from monai.config import print_config
+from monai.data import DataLoader #, Dataset
+from monai.utils import first, set_determinism
+from torch.cuda.amp import GradScaler, autocast
+from pathlib import Path
+# from tqdm import tqdm
+
+from generative.inferers import LatentDiffusionInferer
+# from generative.losses.adversarial_loss import PatchAdversarialLoss
+# from generative.losses.perceptual import PerceptualLoss
+from generative.networks.nets import AutoencoderKL, DiffusionModelUNet # , PatchDiscriminator
+from generative.networks.schedulers import DDPMScheduler
+
+from cycling_utils import InterruptableDistributedSampler, Timer, MetricsTracker
+# from loops import train_generator_one_epoch, evaluate_generator
+from loops import train_diffusion_one_epoch, evaluate_diffusion
+import utils
+
+def get_args_parser(add_help=True):
+    import argparse
+    parser = argparse.ArgumentParser(description="Latent Diffusion Model Training", add_help=add_help)
+    parser.add_argument("--resume", type=str, help="path of checkpoint", required=True) # for checkpointing
+    parser.add_argument("--gen-load-path", type=str, help="path of checkpoint", dest="gen_load_path") # for checkpointing
+    parser.add_argument("--prev-resume", default=None, help="path of previous job checkpoint for strong fail resume", dest="prev_resume") # for checkpointing
+    parser.add_argument("--tboard-path", default=None, help="path for saving tensorboard logs", dest="tboard_path") # for checkpointing
+    parser.add_argument("--data-path", default="/mnt/Datasets/Open-Datasets/MONAI", type=str, help="dataset path", dest="data_path")
+    parser.add_argument("--device", default="cuda", type=str, help="device (Use cuda or cpu Default: cuda)")
+    parser.add_argument("--start-epoch", default=0, type=int, metavar="N", help="start epoch")
+    parser.add_argument("--dist-url", default="env://", type=str, help="url used to set up distributed training")
+    parser.add_argument("-j", "--workers", default=16, type=int, metavar="N", help="number of data loading workers (default: 16)")
+    return parser
+
+def compute_scale_factor(autoencoder, train_loader, device):
+    with torch.no_grad():
+        check_data = first(train_loader)
+        z = autoencoder.encode_stage_2_inputs(check_data["image"].to(device))
+    scale_factor = 1 / torch.std(z)
+    return scale_factor.item()
+
+timer.report('importing everything else')
+
+def main(args, timer):
+
+    utils.init_distributed_mode(args) # Sets args.distributed among other things
+    assert args.distributed # don't support cycling when not distributed for simplicity
+
+    device = torch.device(args.device)
+
+    timer.report('preliminaries')
+
+    # Maybe this will work?
+    set_determinism(42)
+
+    channel = 0  # 0 = "Flair" channel
+    assert channel in [0, 1, 2, 3], "Choose a valid channel"
+    preprocessing_transform = transforms.Compose([
+            transforms.LoadImaged(keys="image", image_only=False), # image_only current default will change soon, so including explicitly
+            transforms.EnsureChannelFirstd(keys="image"),
+            transforms.Lambdad(keys="image", func=lambda x: x[channel, :, :, :]),
+            transforms.AddChanneld(keys="image"),
+            transforms.EnsureTyped(keys="image"),
+            transforms.Orientationd(keys="image", axcodes="RAS"),
+            transforms.CenterSpatialCropd(keys="image", roi_size=(240, 240, 100)),
+            transforms.ScaleIntensityRangePercentilesd(keys="image", lower=0, upper=100, b_min=0, b_max=1),
+    ])
+            
+    crop_transform = transforms.Compose([
+            transforms.DivisiblePadd(keys="image", k=[32,32,1]),
+            transforms.RandSpatialCropd(keys="image", roi_size=(256, 256, 1), random_size=False), # Each of the 100 slices will be randomly sampled.
+            transforms.SqueezeDimd(keys="image", dim=3),
+            # transforms.RandFlipd(keys="image", prob=0.5, spatial_axis=0),
+            # transforms.RandFlipd(keys="image", prob=0.5, spatial_axis=1),
+    ])
+
+    preprocessing = transforms.Compose([preprocessing_transform, crop_transform])
+
+    train_ds = DecathlonDataset(
+        root_dir=args.data_path, task="Task01_BrainTumour", section="training", cache_rate=0.0,
+        num_workers=8, download=False, seed=0, transform=preprocessing,
+    )
+    val_ds = DecathlonDataset(
+        root_dir=args.data_path, task="Task01_BrainTumour", section="validation", cache_rate=0.0,
+        num_workers=8, download=False, seed=0, transform=preprocessing,
+    )
+
+    timer.report('build datasets')
+
+    train_sampler = InterruptableDistributedSampler(train_ds)
+    val_sampler = InterruptableDistributedSampler(val_ds)
+
+    timer.report('build samplers')
+
+    # Original trainer had batch size = 2 * 50. Using 11 nodes x 6 GPUs x batch size 2 => eff batch size = 132
+    train_loader = DataLoader(train_ds, batch_size=2, sampler=train_sampler, num_workers=1)
+    val_loader = DataLoader(val_ds, batch_size=1, sampler=val_sampler, num_workers=1)
+    # check_data = first(train_loader) # Used later
+
+    timer.report('build dataloaders')
+
+    # Auto-encoder definition
+    generator = AutoencoderKL(
+        spatial_dims=2, in_channels=1, out_channels=1, num_channels=(64, 128, 256), 
+        latent_channels=1, num_res_blocks=2, norm_num_groups=32, norm_eps=1e-06,
+        attention_levels=(False, False, False), with_encoder_nonlocal_attn=True, 
+        with_decoder_nonlocal_attn=True,
+    )
+    # saved_generator_checkpoint = torch.load("/output_brats_mri_2d_gen/exp_1645/checkpoint.isc", map_location="cpu")
+    saved_generator_checkpoint = torch.load(args.gen_load_path, map_location="cpu")
+    generator.load_state_dict(saved_generator_checkpoint["generator"])
+    generator = generator.to(device)
+
+    timer.report('generator to device')
+
+    # Diffusion model (unet)
+    unet = DiffusionModelUNet(
+        spatial_dims=2, in_channels=1, out_channels=1, num_res_blocks=2, 
+        num_channels=(32, 64, 128, 256), attention_levels=(False, True, True, True), 
+        num_head_channels=(0, 32, 32, 32),
+    )
+    unet = unet.to(device)
+
+    timer.report('unet to device')
+
+    # Prepare for distributed training
+    unet = torch.nn.SyncBatchNorm.convert_sync_batchnorm(unet)
+
+    unet_without_ddp = unet
+    if args.distributed:
+        unet = torch.nn.parallel.DistributedDataParallel(unet, device_ids=[args.gpu], find_unused_parameters=True)
+        unet_without_ddp = unet.module
+
+    timer.report('unet prepped for distribution')
+
+    # Optimizers
+    optimizer_u = torch.optim.Adam(unet_without_ddp.parameters(), lr=5e-5)
+    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer_u, milestones=[1000], gamma=0.1)
+
+    # For mixed precision training
+    scaler_u = GradScaler()
+
+    timer.report('optimizer, lr_scheduler and grad scaler')
+
+    # Init metric tracker
+    metrics = {'train': MetricsTracker(), 'val': MetricsTracker()}
+
+    # Prepare LatentDiffusionInferer
+
+    scale_factor = compute_scale_factor(generator, train_loader, device)
+    scheduler = DDPMScheduler(num_train_timesteps=1000, schedule="scaled_linear_beta", beta_start=0.0015, beta_end=0.0195)
+    inferer = LatentDiffusionInferer(scheduler, scale_factor=scale_factor)
+
+    timer.report('building inferer')
+    
+    # RETRIEVE CHECKPOINT
+    Path(args.resume).parent.mkdir(parents=True, exist_ok=True)
+    checkpoint = None
+    if args.resume and os.path.isfile(args.resume): # If we're resuming...
+        checkpoint = torch.load(args.resume, map_location="cpu")
+    elif args.prev_resume and os.path.isfile(args.prev_resume):
+        checkpoint = torch.load(args.prev_resume, map_location="cpu")
+    if checkpoint is not None:
+        args.start_epoch = checkpoint["epoch"]
+        unet_without_ddp.load_state_dict(checkpoint["unet"])
+        optimizer_u.load_state_dict(checkpoint["optimizer_u"])
+        scaler_u.load_state_dict(checkpoint["scaler_u"])
+        train_sampler.load_state_dict(checkpoint["train_sampler"])
+        val_sampler.load_state_dict(checkpoint["val_sampler"])
+        lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
+        # Metrics
+        metrics = checkpoint["metrics"]
+
+    timer.report('checkpoint retrieval')
+
+    ## -- TRAINING THE DIFFUSION MODEL - ##
+
+    n_diff_epochs = 200
+    diff_val_interval = 1
+
+    for epoch in range(args.start_epoch, n_diff_epochs):
+
+        print('\n')
+        print(f"EPOCH :: {epoch}")
+        print('\n')
+
+        with train_sampler.in_epoch(epoch):
+            timer = TimestampedTimer("Start training")
+            unet, timer, metrics = train_diffusion_one_epoch(
+                args, epoch, unet, generator, optimizer_u, scaler_u, inferer, train_loader, val_loader, 
+                train_sampler, val_sampler, lr_scheduler, device, timer, metrics
+            )
+            timer.report(f'training unet for epoch {epoch}')
+
+            if epoch % diff_val_interval == 0:
+                with val_sampler.in_epoch(epoch):
+                    timer = TimestampedTimer("Start evaluation")
+                    timer, metrics = evaluate_diffusion(
+                        args, epoch, unet, generator, optimizer_u, scaler_u, inferer, train_loader, val_loader, 
+                        train_sampler, val_sampler, lr_scheduler, device, timer, metrics
+                    )
+                    timer.report(f'evaluating unet for epoch {epoch}')
+
+
+if __name__ == "__main__":
+    args = get_args_parser().parse_args()
+    main(args, timer)
diff --git a/monai_brats_mri_2d/train_cycling_gen.py b/monai_brats_mri_2d/train_cycling_gen.py
new file mode 100644
index 00000000..8240e853
--- /dev/null
+++ b/monai_brats_mri_2d/train_cycling_gen.py
@@ -0,0 +1,210 @@
+from cycling_utils import TimestampedTimer
+
+timer = TimestampedTimer()
+timer.report('importing Timer')
+
+import torch, os, utils
+from torch.cuda.amp import GradScaler
+from pathlib import Path
+# !pip install monai==1.2.0
+from monai import transforms
+from monai.apps import DecathlonDataset
+from monai.data import DataLoader
+from monai.utils import set_determinism
+from generative.losses.perceptual import PerceptualLoss
+from generative.networks.nets import AutoencoderKL, PatchDiscriminator
+
+from cycling_utils import InterruptableDistributedSampler, MetricsTracker
+from loops import train_generator_one_epoch, evaluate_generator
+
+def get_args_parser(add_help=True):
+    import argparse
+    parser = argparse.ArgumentParser(description="Latent Diffusion Model Training", add_help=add_help)
+    parser.add_argument("--resume", type=str, help="path of checkpoint", required=True) # for checkpointing
+    parser.add_argument("--prev-resume", default=None, help="path of previous job checkpoint for strong fail resume", dest="prev_resume") # for checkpointing
+    parser.add_argument("--tboard-path", default=None, help="path for saving tensorboard logs", dest="tboard_path") # for checkpointing
+    parser.add_argument("--data-path", default="/mnt/Datasets/Open-Datasets/MONAI", type=str, help="dataset path", dest="data_path")
+    parser.add_argument("--device", default="cuda", type=str, help="device (Use cuda or cpu Default: cuda)")
+    parser.add_argument("--start-epoch", default=0, type=int, metavar="N", help="start epoch")
+    parser.add_argument("--dist-url", default="env://", type=str, help="url used to set up distributed training")
+    parser.add_argument("-j", "--workers", default=16, type=int, metavar="N", help="number of data loading workers (default: 16)")
+    parser.add_argument("--lr", default=5e-5, type=float, help="initial learning rate")
+    parser.add_argument("--kl-weight",default=1e-6,type=float, help="kl loss weight for generator", dest="kl_weight")
+    parser.add_argument("--perceptual-weight",default=1.0,type=float, help="perceptual loss weight for generator", dest="perceptual_weight")
+    parser.add_argument("--adv-weight",default=0.5,type=float, help="adversarial loss weight for generator", dest="adv_weight")
+    return parser
+
+timer.report('importing everything else')
+
+def main(args, timer):
+
+    utils.init_distributed_mode(args) # Sets args.distributed among other things
+    assert args.distributed # don't support cycling when not distributed for simplicity
+    device = torch.device(args.device)
+    set_determinism(42)
+
+    timer.report('preliminaries')
+
+    channel = 0  # 0 = "Flair" channel
+    assert channel in [0, 1, 2, 3], "Choose a valid channel"
+    preprocessing_transform = transforms.Compose([
+            transforms.LoadImaged(keys="image", image_only=False), # image_only current default will change soon, so including explicitly
+            transforms.EnsureChannelFirstd(keys="image"),
+            transforms.Lambdad(keys="image", func=lambda x: x[channel, :, :, :]),
+            transforms.AddChanneld(keys="image"),
+            transforms.EnsureTyped(keys="image"),
+            transforms.Orientationd(keys="image", axcodes="RAS"),
+            transforms.CenterSpatialCropd(keys="image", roi_size=(240, 240, 100)),
+            transforms.ScaleIntensityRangePercentilesd(keys="image", lower=0, upper=100, b_min=0, b_max=1),
+    ])
+            
+    crop_transform = transforms.Compose([
+            transforms.DivisiblePadd(keys="image", k=[4,4,1]),
+            transforms.RandSpatialCropd(keys="image", roi_size=(240, 240, 1), random_size=False), # Each of the 100 slices will be randomly sampled.
+            # transforms.RandSpatialCropSamplesd(keys="image", random_size=False, roi_size=(240, 240, 1), num_samples=26), # Each of the 100 slices will be randomly sampled.
+            transforms.SqueezeDimd(keys="image", dim=3),
+            transforms.RandFlipd(keys="image", prob=0.5, spatial_axis=0),
+            transforms.RandFlipd(keys="image", prob=0.5, spatial_axis=1),
+    ])
+
+    preprocessing = transforms.Compose([preprocessing_transform, crop_transform])
+
+    train_ds = DecathlonDataset(
+        root_dir=args.data_path, task="Task01_BrainTumour", section="training", cache_rate=0.0,
+        num_workers=8, download=False, seed=0, transform=preprocessing,
+    )
+    val_ds = DecathlonDataset(
+        root_dir=args.data_path, task="Task01_BrainTumour", section="validation", cache_rate=0.0,
+        num_workers=8, download=False, seed=0, transform=preprocessing,
+    )
+
+    timer.report('build datasets')
+
+    train_sampler = InterruptableDistributedSampler(train_ds)
+    val_sampler = InterruptableDistributedSampler(val_ds)
+
+    timer.report('build samplers')
+
+    # Original trainer had batch size = 26. Using 11 nodes x 6 GPUs x batch size 1 = eff batch size = 66
+    train_loader = DataLoader(train_ds, batch_size=1, sampler=train_sampler, num_workers=1)
+    val_loader = DataLoader(val_ds, batch_size=1, sampler=val_sampler, num_workers=1)
+
+    timer.report('build dataloaders')
+
+    # Auto-encoder definition
+    generator = AutoencoderKL(
+        spatial_dims=2, in_channels=1, out_channels=1, num_channels=(64, 128, 256), 
+        latent_channels=1, num_res_blocks=2, norm_num_groups=32, norm_eps=1e-06,
+        attention_levels=(False, False, False), with_encoder_nonlocal_attn=True, 
+        with_decoder_nonlocal_attn=True,
+    )
+    generator = generator.to(device)
+
+    timer.report('generator to device')
+
+    # Discriminator definition
+    discriminator = PatchDiscriminator(
+        spatial_dims=2, num_layers_d=3, num_channels=32, 
+        in_channels=1, out_channels=1, norm="INSTANCE"
+    )
+    discriminator = discriminator.to(device)
+
+    timer.report('discriminator to device')
+
+    # Autoencoder loss functions
+    perceptual_loss = PerceptualLoss(
+        spatial_dims=2, network_type="resnet50", pretrained=True, #ImageNet pretrained weights used
+    )
+    perceptual_loss.to(device)
+
+    timer.report('loss functions')
+
+    # Prepare for distributed training
+    generator = torch.nn.SyncBatchNorm.convert_sync_batchnorm(generator)
+    discriminator = torch.nn.SyncBatchNorm.convert_sync_batchnorm(generator)
+
+    generator_without_ddp = generator
+    discriminator_without_ddp = discriminator
+    if args.distributed:
+        generator = torch.nn.parallel.DistributedDataParallel(generator, device_ids=[args.gpu], find_unused_parameters=True) # find_unused_parameters necessary for monai training
+        discriminator = torch.nn.parallel.DistributedDataParallel(discriminator, device_ids=[args.gpu], find_unused_parameters=True) # find_unused_parameters necessary for monai training
+        generator_without_ddp = generator.module
+        discriminator_without_ddp = discriminator.module
+
+    timer.report('models prepped for distribution')
+
+    # Optimizers
+    optimizer_g = torch.optim.Adam(generator_without_ddp.parameters(), lr=args.lr)
+    optimizer_d = torch.optim.Adam(discriminator_without_ddp.parameters(), lr=args.lr)
+
+    timer.report('optimizers')
+
+    # For mixed precision training
+    scaler_g = GradScaler()
+    scaler_d = GradScaler()
+
+    timer.report('grad scalers')
+
+    # Init metric tracker
+    metrics = {'train': MetricsTracker(), 'val': MetricsTracker()}
+
+    # RETRIEVE CHECKPOINT
+    Path(args.resume).parent.mkdir(parents=True, exist_ok=True)
+    checkpoint = None
+    if args.resume and os.path.isfile(args.resume): # If we're resuming...
+        checkpoint = torch.load(args.resume, map_location="cpu")
+        print("RESUMING PAUSED EXPERIMENT")
+    elif args.prev_resume and os.path.isfile(args.prev_resume):
+        checkpoint = torch.load(args.prev_resume, map_location="cpu")
+        print("RE-STARTING FROM PREVIOUS EXPERIMENT")
+
+    if checkpoint is not None:
+        args.start_epoch = checkpoint["epoch"]
+        generator_without_ddp.load_state_dict(checkpoint["generator"])
+        discriminator_without_ddp.load_state_dict(checkpoint["discriminator"])
+        optimizer_g.load_state_dict(checkpoint["optimizer_g"])
+        optimizer_d.load_state_dict(checkpoint["optimizer_d"])
+        scaler_g.load_state_dict(checkpoint["scaler_g"])
+        scaler_d.load_state_dict(checkpoint["scaler_d"])
+        train_sampler.load_state_dict(checkpoint["train_sampler"])
+        val_sampler.load_state_dict(checkpoint["val_sampler"])
+        # Metrics
+        metrics = checkpoint["metrics"]
+
+    timer.report('checkpoint retrieval')
+
+    ## -- TRAINING THE AUTO-ENCODER - ##
+
+    n_gen_epochs = 100_000
+    gen_val_interval = 1
+
+    for epoch in range(args.start_epoch, n_gen_epochs):
+
+        print('\n')
+        print(f"EPOCH :: {epoch}")
+        print('\n')
+
+        with train_sampler.in_epoch(epoch):
+            timer = TimestampedTimer("Start training")
+
+            generator, timer, metrics = train_generator_one_epoch(
+                args, epoch, generator, discriminator, optimizer_g, optimizer_d, train_sampler, val_sampler,
+                scaler_g, scaler_d, train_loader, perceptual_loss, device, timer, metrics
+            )
+            timer.report(f'training generator for epoch {epoch}')
+
+            if epoch % gen_val_interval == 0: # Eval every epoch
+                
+                with val_sampler.in_epoch(epoch):
+                    timer = TimestampedTimer("Start evaluation")
+
+                    timer, metrics = evaluate_generator(
+                        args, epoch, generator, discriminator, optimizer_g, optimizer_d, train_sampler, val_sampler,
+                        scaler_g, scaler_d, val_loader, device, timer, metrics
+                    )
+                    timer.report(f'evaluating generator for epoch {epoch}')
+
+
+if __name__ == "__main__":
+    args = get_args_parser().parse_args()
+    main(args, timer)
diff --git a/monai_brats_mri_2d/utils.py b/monai_brats_mri_2d/utils.py
new file mode 100644
index 00000000..1b268b82
--- /dev/null
+++ b/monai_brats_mri_2d/utils.py
@@ -0,0 +1,72 @@
+import torch, os, errno
+import torch.distributed as dist
+
+def mkdir(path):
+    try:
+        os.makedirs(path)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+def init_distributed_mode(args):
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print("Not using distributed mode")
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = "nccl"
+    print(f"| distributed init (rank {args.rank}): {args.dist_url}", flush=True)
+    torch.distributed.init_process_group(
+        backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank
+    )
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
\ No newline at end of file
diff --git a/monai_brats_mri_3d/requirements-dev.txt b/monai_brats_mri_3d/requirements-dev.txt
new file mode 100644
index 00000000..c772949e
--- /dev/null
+++ b/monai_brats_mri_3d/requirements-dev.txt
@@ -0,0 +1,57 @@
+# Full requirements for developments
+-r requirements-min.txt
+pytorch-ignite==0.4.10
+gdown>=4.4.0
+scipy
+itk>=5.2
+nibabel
+pillow!=8.3.0  # https://github.com/python-pillow/Pillow/issues/5571
+tensorboard>=2.6  # https://github.com/Project-MONAI/MONAI/issues/5776
+scikit-image>=0.19.0
+tqdm>=4.47.0
+lmdb
+flake8>=3.8.1
+flake8-bugbear
+flake8-comprehensions
+flake8-executable
+pylint!=2.13  # https://github.com/PyCQA/pylint/issues/5969
+mccabe
+pep8-naming
+pycodestyle
+pyflakes
+black
+isort
+pytype>=2020.6.1; platform_system != "Windows"
+types-pkg_resources
+mypy>=0.790
+ninja
+# torchvision
+psutil
+Sphinx==3.5.3
+recommonmark==0.6.0
+sphinx-autodoc-typehints==1.11.1
+sphinx-rtd-theme==0.5.2
+cucim==22.8.1; platform_system == "Linux"
+openslide-python==1.1.2
+imagecodecs; platform_system == "Linux" or platform_system == "Darwin"
+tifffile; platform_system == "Linux" or platform_system == "Darwin"
+pandas
+requests
+einops
+transformers<4.22  # https://github.com/Project-MONAI/MONAI/issues/5157
+mlflow
+matplotlib!=3.5.0
+tensorboardX
+types-PyYAML
+pyyaml
+fire
+jsonschema
+pynrrd
+pre-commit
+pydicom
+h5py
+nni
+optuna
+git+https://github.com/Project-MONAI/MetricsReloaded@monai-support#egg=MetricsReloaded
+lpips==0.1.4
+xformers==0.0.16
diff --git a/monai_brats_mri_3d/requirements-min.txt b/monai_brats_mri_3d/requirements-min.txt
new file mode 100644
index 00000000..ceb9e346
--- /dev/null
+++ b/monai_brats_mri_3d/requirements-min.txt
@@ -0,0 +1,5 @@
+# Requirements for minimal tests
+-r requirements.txt
+setuptools>65.5.0,<66.0.0
+coverage>=5.5
+parameterized
diff --git a/monai_brats_mri_3d/requirements.txt b/monai_brats_mri_3d/requirements.txt
new file mode 100644
index 00000000..3f1ff86c
--- /dev/null
+++ b/monai_brats_mri_3d/requirements.txt
@@ -0,0 +1,3 @@
+numpy>=1.17
+torch>=1.8
+monai>=1.2.0rc1
diff --git a/monai_pancreas_dints/LICENSE b/monai_pancreas_dints/LICENSE
new file mode 100644
index 00000000..261eeb9e
--- /dev/null
+++ b/monai_pancreas_dints/LICENSE
@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
diff --git a/monai_pancreas_dints/SS/loops_SS.py b/monai_pancreas_dints/SS/loops_SS.py
new file mode 100644
index 00000000..db5a8db7
--- /dev/null
+++ b/monai_pancreas_dints/SS/loops_SS.py
@@ -0,0 +1,497 @@
+
+import torch
+import numpy as np
+import torch.distributed as dist
+from torch.cuda.amp import autocast
+# from datetime import datetime
+# from scipy import ndimage
+import torch.nn.functional as F
+
+from monai.inferers import sliding_window_inference
+from monai.metrics import compute_dice
+# import yaml, time, 
+import os
+import utils
+from cycling_utils import atomic_torch_save
+from torch.utils.tensorboard import SummaryWriter
+
+def search_one_epoch(
+    model, optimizer, dints_space, arch_optimizer_a, arch_optimizer_c, 
+    train_sampler, val_sampler, scaler, train_metrics, val_metric,
+    epoch, train_loader, loss_func, args, timer
+):
+    device = args["device"] # for convenience
+
+    decay = 0.5 ** np.sum(
+        [(epoch - args["num_epochs_warmup"]) / (args["num_epochs"] - args["num_epochs_warmup"]) > args["learning_rate_milestones"]]
+    )
+    lr = args["learning_rate"] * decay * args["world_size"]
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr
+
+    model.train()
+
+    timer.report('model.train()')
+
+    train_step = train_sampler.progress // train_loader.batch_size
+    total_steps = int(len(train_sampler) / train_loader.batch_size)
+    print(f'\nTraining / resuming epoch {epoch} from training step {train_step}\n')
+
+    for batch_data in train_loader:
+
+        inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
+        inputs_search, labels_search = inputs.detach().clone(), labels.detach().clone() # added, will this work?
+
+        timer.report('data to device')
+
+        # UPDATE MODEL
+
+        for p in model.module.weight_parameters():
+            p.requires_grad=True
+        dints_space.log_alpha_a.requires_grad = False
+        dints_space.log_alpha_c.requires_grad = False
+
+        optimizer.zero_grad()
+
+        timer.report('config model to train')
+
+        if args["amp"]:
+            with autocast():
+                outputs = model(inputs)
+                timer.report('model forward')
+                if args["output_classes"] == 2:
+                    loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+                else:
+                    loss = loss_func(outputs, labels)
+                timer.report('model loss')
+
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+            timer.report('model backward')
+        else:
+            outputs = model(inputs)
+            timer.report('model forward')
+            if args["output_classes"] == 2:
+                loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+            else:
+                loss = loss_func(outputs, labels)
+            timer.report('model loss')
+            loss.backward()
+            optimizer.step()
+            timer.report('model backward')
+
+        # Reporting and stuff
+        train_metrics.update({"model_loss": loss.item(), "inputs_seen": len(inputs)})
+
+        timer.report('model update')
+
+        # Only update space after number of warmup epochs
+        if epoch >= args["num_epochs_warmup"]:
+
+            # UPDATE SPACE
+
+            for p in model.module.weight_parameters():
+                p.requires_grad=False
+            dints_space.log_alpha_a.requires_grad = True
+            dints_space.log_alpha_c.requires_grad = True
+
+            # linear increase topology and RAM loss
+            entropy_alpha_c = torch.tensor(0.0,).to(device)
+            entropy_alpha_a = torch.tensor(0.0).to(device)
+            ram_cost_full = torch.tensor(0.0).to(device)
+            ram_cost_usage = torch.tensor(0.0).to(device)
+            ram_cost_loss = torch.tensor(0.0).to(device)
+            topology_loss = torch.tensor(0.0).to(device)
+
+            probs_a, arch_code_prob_a = dints_space.get_prob_a(child=True)
+            entropy_alpha_a = -((probs_a) * torch.log(probs_a + 1e-5)).mean()
+            sm = F.softmax(dints_space.log_alpha_c, dim=-1)
+            lsm = F.log_softmax(dints_space.log_alpha_c, dim=-1)
+            entropy_alpha_c = -(sm * lsm).mean()
+            topology_loss = dints_space.get_topology_entropy(probs_a)
+
+            ram_cost_full = dints_space.get_ram_cost_usage(inputs.shape, full=True)
+            ram_cost_usage = dints_space.get_ram_cost_usage(inputs.shape)
+            ram_cost_loss = torch.abs(args["ram_cost_factor"] - ram_cost_usage / ram_cost_full)
+
+            arch_optimizer_a.zero_grad()
+            arch_optimizer_c.zero_grad()
+
+            combination_weights = (epoch - args["num_epochs_warmup"]) / (args["num_epochs"] - args["num_epochs_warmup"])
+
+            timer.report('space combination_weights')
+
+            if args["amp"]:
+                with autocast():
+                    outputs_search = model(inputs_search)
+                    timer.report('space forward')
+                    if args["output_classes"] == 2:
+                        loss = loss_func(torch.flip(outputs_search, dims=[1]), 1 - labels_search)
+                    else:
+                        loss = loss_func(outputs_search, labels_search)
+
+                    loss += combination_weights * (
+                        (entropy_alpha_a + entropy_alpha_c) + ram_cost_loss + 0.001 * topology_loss
+                    )
+                    timer.report('space loss')
+
+                scaler.scale(loss).backward()
+                scaler.step(arch_optimizer_a)
+                scaler.step(arch_optimizer_c)
+                scaler.update()
+                timer.report('space backward')
+            else:
+                outputs_search = model(inputs_search)
+                timer.report('space forward')
+                if args["output_classes"] == 2:
+                    loss = loss_func(torch.flip(outputs_search, dims=[1]), 1 - labels_search)
+                else:
+                    loss = loss_func(outputs_search, labels_search)
+
+                loss += 1.0 * (
+                    combination_weights * (entropy_alpha_a + entropy_alpha_c) + ram_cost_loss + 0.001 * topology_loss
+                )
+                timer.report('space loss')
+
+                loss.backward()
+                arch_optimizer_a.step()
+                arch_optimizer_c.step()
+                timer.report('space backward')
+
+            # Reporting and stuff
+            train_metrics.update({"space_loss": loss.item()})
+
+            timer.report('space update')
+
+        # Batch reporting
+        train_metrics.reduce()
+        batch_model_loss = train_metrics.local["model_loss"] / train_metrics.local["inputs_seen"]
+        if "space_loss" in train_metrics.local:
+            batch_space_loss = train_metrics.local["space_loss"] / train_metrics.local["inputs_seen"]
+        else:
+            batch_space_loss = 0.0
+        print(f"EPOCH [{epoch}], BATCH [{train_step}], MODEL LOSS [{batch_model_loss:,.3f}, SPACE LOSS: [{batch_space_loss:,.3f}]")
+        train_metrics.reset_local()
+
+        timer.report('metrics reduce')
+
+        ## Checkpointing
+        print(f"Saving checkpoint at epoch {epoch} train batch {train_step}")
+        train_sampler.advance(len(inputs))
+        train_step = train_sampler.progress // train_loader.batch_size
+
+        if train_step == total_steps:
+            train_metrics.end_epoch()
+
+        if utils.is_main_process() and train_step % 1 == 0: # Checkpointing every batch
+
+            writer = SummaryWriter(log_dir=args["tboard_path"])
+            writer.add_scalar("Train/model_loss", batch_model_loss, train_step + epoch * total_steps)
+            if batch_space_loss != "NONE":
+                writer.add_scalar("Train/space_loss", batch_space_loss, train_step + epoch * total_steps)
+            writer.flush()
+            writer.close()
+
+            checkpoint = {
+                "epoch": epoch,
+                "model": model.module.state_dict(),
+                "dints": dints_space.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "arch_optimizer_a": arch_optimizer_a.state_dict(),
+                "arch_optimizer_c": arch_optimizer_c.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "val_sampler": val_sampler.state_dict(),
+                "scaler": scaler.state_dict(),
+                "train_metrics": train_metrics,
+                "val_metric": val_metric
+            }
+            timer = atomic_torch_save(checkpoint, args["resume"], timer)
+
+    timer.report(f'EPOCH {epoch}')
+
+    return model, dints_space, timer, train_metrics
+
+
+def eval_search(
+        model, optimizer, dints_space, arch_optimizer_a, arch_optimizer_c, 
+        train_sampler, val_sampler, scaler, train_metrics, val_metric,
+        epoch, val_loader, post_pred, post_label, args, timer,
+):
+    device = args["device"] # for convenience
+
+    torch.cuda.empty_cache()
+    model.eval()
+
+    timer.report('model ready to eval')
+
+    with torch.no_grad():
+
+        val_step = val_sampler.progress // val_loader.batch_size
+        total_steps = int(len(val_sampler) / val_loader.batch_size)
+        print(f'\nEvaluating / resuming epoch {epoch} from eval step {val_step}\n')
+
+        for val_data in val_loader:
+
+            val_images, val_labels = val_data["image"].to(device), val_data["label"].to(device)
+            roi_size = args["patch_size_valid"]
+            sw_batch_size = args["num_sw_batch_size"]
+
+            if args["amp"]:
+                with torch.cuda.amp.autocast():
+                    pred = sliding_window_inference(
+                        val_images, roi_size, sw_batch_size,
+                        lambda x: model(x), mode="gaussian",
+                        overlap=args["overlap_ratio"],
+                    )
+            else:
+                pred = sliding_window_inference(
+                    val_images, roi_size, sw_batch_size,
+                    lambda x: model(x), mode="gaussian",
+                    overlap=args["overlap_ratio"],
+                )
+
+            val_outputs = post_pred(pred[0, ...])
+            val_outputs = val_outputs[None, ...]
+            val_labels = post_label(val_labels[0, ...])
+            val_labels = val_labels[None, ...]
+
+            value = compute_dice(y_pred=val_outputs, y=val_labels, include_background=False)
+
+            for _c in range(args["output_classes"] - 1):
+                val0 = torch.nan_to_num(value[0, _c], nan=0.0)
+                val1 = 1.0 - torch.isnan(value[0, 0]).float()
+                val_metric[2 * _c] += val0 * val1
+                val_metric[2 * _c + 1] += val1
+
+            ## Checkpointing
+            print(f"Saving checkpoint at epoch {epoch} eval batch {val_step}")
+            val_sampler.advance(len(val_images))
+            val_step = val_sampler.progress // val_loader.batch_size
+
+            if utils.is_main_process() and val_step % 1 == 0: # Checkpointing every batch
+
+                checkpoint = {
+                    "epoch": epoch,
+                    "model": model.module.state_dict(),
+                    "dints": dints_space.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "arch_optimizer_a": arch_optimizer_a.state_dict(),
+                    "arch_optimizer_c": arch_optimizer_c.state_dict(),
+                    "train_sampler": train_sampler.state_dict(),
+                    "val_sampler": val_sampler.state_dict(),
+                    "scaler": scaler.state_dict(),
+                    "train_metrics": train_metrics,
+                    "val_metric": val_metric
+                }
+                timer = atomic_torch_save(checkpoint, args["resume"], timer)
+
+            timer.report(f'eval step {val_step}')
+
+        # synchronizes all processes and reduce results
+        if torch.cuda.device_count() > 1:
+            dist.barrier()
+            dist.all_reduce(val_metric, op=torch.distributed.ReduceOp.SUM)
+
+        val_metric = val_metric.tolist()
+        if utils.is_main_process():
+
+            for _c in range(args["output_classes"] - 1):
+                print("evaluation metric - class {0:d}:".format(_c + 1), val_metric[2 * _c] / val_metric[2 * _c + 1])
+            avg_metric = 0
+            for _c in range(args["output_classes"] - 1):
+                avg_metric += val_metric[2 * _c] / val_metric[2 * _c + 1]
+            avg_metric = avg_metric / float(args["output_classes"] - 1)
+            print("avg_metric", avg_metric)
+
+            if avg_metric > best_metric:
+                best_metric = avg_metric
+                # best_metric_epoch = epoch + 1
+                # best_metric_iterations = idx_iter
+
+                (node_a_d, arch_code_a_d, arch_code_c_d, arch_code_a_max_d) = dints_space.decode()
+                torch.save(
+                    {
+                        "node_a": node_a_d,
+                        "arch_code_a": arch_code_a_d,
+                        "arch_code_a_max": arch_code_a_max_d,
+                        "arch_code_c": arch_code_c_d,
+                        # "iter_num": idx_iter,
+                        "epochs": epoch + 1,
+                        "best_dsc": best_metric,
+                        # "best_path": best_metric_iterations,
+                    },
+                    os.path.join(args["arch_ckpt_path"], "search_code.pt"),
+                )
+    
+    timer.report(f'EVAL EPOCH {epoch}')
+
+    return timer
+
+
+def train_one_epoch(
+    model, optimizer, lr_scheduler,
+    train_sampler, val_sampler, scaler, train_metrics, val_metric,
+    epoch, train_loader, loss_func, args
+):
+    device = args["device"] # for convenience
+
+    model.train()
+
+    train_step = train_sampler.progress // train_loader.batch_size
+    total_steps = int(len(train_sampler) / train_loader.batch_size)
+    print(f'\nTraining / resuming epoch {epoch} from training step {train_step}\n')
+
+    for batch_data in train_loader:
+
+        inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
+
+        optimizer.zero_grad()
+
+        if args["amp"]:
+            with autocast():
+                outputs = model(inputs)
+                if args["output_classes"] == 2:
+                    loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+                else:
+                    loss = loss_func(outputs, labels)
+
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            outputs = model(inputs)
+            if args["output_classes"] == 2:
+                loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+            else:
+                loss = loss_func(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+        # Reporting and stuff
+        train_metrics.update({"model_loss": loss.item(), "inputs_seen": len(inputs)})
+        train_metrics.reduce()
+        batch_model_loss = train_metrics.local["model_loss"] / train_metrics.local["inputs_seen"]
+        print(f"EPOCH [{epoch}], BATCH [{train_step}], MODEL LOSS [{batch_model_loss:,.3f}]")
+        train_metrics.reset_local()
+
+        ## Checkpointing
+        print(f"Saving checkpoint at epoch {epoch} train batch {train_step}")
+        train_sampler.advance(len(inputs))
+        train_step = train_sampler.progress // train_loader.batch_size
+
+        if train_step == total_steps:
+            train_metrics.end_epoch()
+            lr_scheduler.step()
+
+        if utils.is_main_process() and train_step % 1 == 0: # Checkpointing every batch
+
+            writer = SummaryWriter(log_dir=args["tboard_path"])
+            writer.add_scalar("Train/model_loss", batch_model_loss, train_step + epoch * total_steps)
+            writer.flush()
+            writer.close()
+
+            checkpoint = {
+                "epoch": epoch,
+                "model": model.module.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "val_sampler": val_sampler.state_dict(),
+                "scaler": scaler.state_dict(),
+                "train_metrics": train_metrics,
+                "val_metric": val_metric,
+            }
+            timer = atomic_torch_save(checkpoint, args["resume"], timer)
+
+    return model, timer, train_metrics
+
+
+def evaluate(
+        model, optimizer, lr_scheduler,
+        train_sampler, val_sampler, scaler, train_metrics, val_metric,
+        epoch, val_loader, post_pred, post_label, args,
+):
+    device = args["device"] # for convenience
+
+    torch.cuda.empty_cache()
+    model.eval()
+
+    with torch.no_grad():
+
+        val_step = val_sampler.progress // val_loader.batch_size
+        print(f'\nEvaluating / resuming epoch {epoch} from eval step {val_step}\n')
+
+        for val_data in val_loader:
+
+            val_images, val_labels = val_data["image"].to(device), val_data["label"].to(device)
+            roi_size = args["patch_size_valid"]
+            sw_batch_size = args["num_sw_batch_size"]
+
+            if args["amp"]:
+                with torch.cuda.amp.autocast():
+                    pred = sliding_window_inference(
+                        val_images, roi_size, sw_batch_size,
+                        lambda x: model(x), mode="gaussian",
+                        overlap=args["overlap_ratio"],
+                    )
+            else:
+                pred = sliding_window_inference(
+                    val_images, roi_size, sw_batch_size,
+                    lambda x: model(x), mode="gaussian",
+                    overlap=args["overlap_ratio"],
+                )
+
+            val_outputs = post_pred(pred[0, ...])
+            val_outputs = val_outputs[None, ...]
+            val_labels = post_label(val_labels[0, ...])
+            val_labels = val_labels[None, ...]
+
+            value = compute_dice(y_pred=val_outputs, y=val_labels, include_background=False)
+
+            for _c in range(args["output_classes"] - 1):
+                val0 = torch.nan_to_num(value[0, _c], nan=0.0)
+                val1 = 1.0 - torch.isnan(value[0, 0]).float()
+                val_metric[2 * _c] += val0 * val1
+                val_metric[2 * _c + 1] += val1
+
+            ## Checkpointing
+            print(f"Saving checkpoint at epoch {epoch} eval batch {val_step}")
+            val_sampler.advance(len(val_images))
+            val_step = val_sampler.progress // val_loader.batch_size
+
+            if utils.is_main_process() and val_step % 1 == 0: # Checkpointing every batch
+
+                checkpoint = {
+                    "epoch": epoch,
+                    "model": model.module.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "lr_scheduler": lr_scheduler.state_dict(),
+                    "train_sampler": train_sampler.state_dict(),
+                    "val_sampler": val_sampler.state_dict(),
+                    "scaler": scaler.state_dict(),
+                    "train_metrics": train_metrics,
+                    "val_metric": val_metric,
+                }
+                timer = atomic_torch_save(checkpoint, args["resume"], timer)
+
+        # synchronizes all processes and reduce results
+        if torch.cuda.device_count() > 1:
+            dist.barrier()
+            dist.all_reduce(val_metric, op=torch.distributed.ReduceOp.SUM)
+
+        val_metric = val_metric.tolist()
+        if utils.is_main_process():
+
+            for _c in range(args["output_classes"] - 1):
+                print("evaluation metric - class {0:d}:".format(_c + 1), val_metric[2 * _c] / val_metric[2 * _c + 1])
+            avg_metric = 0
+            for _c in range(args["output_classes"] - 1):
+                avg_metric += val_metric[2 * _c] / val_metric[2 * _c + 1]
+            avg_metric = avg_metric / float(args["output_classes"] - 1)
+            print("avg_metric", avg_metric)
+
+            writer = SummaryWriter(log_dir=args["tboard_path"])
+            writer.add_scalar("Val/avg_metric", avg_metric, epoch)
+            writer.flush()
+            writer.close()
\ No newline at end of file
diff --git a/monai_pancreas_dints/SS/search_SS.py b/monai_pancreas_dints/SS/search_SS.py
new file mode 100644
index 00000000..04b3794b
--- /dev/null
+++ b/monai_pancreas_dints/SS/search_SS.py
@@ -0,0 +1,264 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from cycling_utils import TimestampedTimer
+
+timer = TimestampedTimer('importing timer')
+
+import json
+# import logging
+import os
+import random
+# import sys
+# import time
+# from datetime import datetime
+# from typing import Sequence, Union
+# from scipy import ndimage
+
+import monai
+import numpy as np
+import torch
+import torch.distributed as dist
+# import torch.nn.functional as F
+# import yaml
+from monai import transforms
+from monai.bundle import ConfigParser
+from monai.networks.nets import TopologySearch, DiNTS
+from monai.losses import DiceCELoss
+# from monai.data import ThreadDataLoader, partition_dataset, 
+from monai.data import DataLoader
+# from monai.inferers import sliding_window_inference
+# from monai.metrics import compute_dice
+from monai.utils import set_determinism
+from torch.nn.parallel import DistributedDataParallel
+
+from cycling_utils import InterruptableDistributedSampler, MetricsTracker
+from loops import search_one_epoch, eval_search
+from pathlib import Path
+import utils
+
+def get_args_parser(add_help=True):
+    import argparse
+    parser = argparse.ArgumentParser(description="DiNTS search", add_help=add_help)
+    parser.add_argument("--config-file", type=str, help="config file", required=True, dest="config_file") # for checkpointing
+    parser.add_argument("--resume", type=str, help="path of checkpoint", required=True) # for checkpointing
+    parser.add_argument("--prev-resume", default=None, help="path of previous job checkpoint for strong fail resume", dest="prev_resume") # for checkpointing
+    parser.add_argument("--tboard-path", default=None, help="path for saving tensorboard logs", dest="tboard_path") # for checkpointing
+    return parser
+
+timer.report('importing everything else')
+
+def main(args, timer):
+    # logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+    timer = TimestampedTimer('commencing run')
+
+    parser = ConfigParser()
+    parser.read_config(args.config_file)
+
+    args = {
+        "start_epoch": 0,
+        "resume": args.resume,
+        "prev_resume": args.prev_resume,
+        "tboard_path": args.tboard_path,
+        "device": "cuda",
+        "dist_url": "env://",
+        "arch_ckpt_path": parser["arch_ckpt_path"],
+        "amp": parser["amp"],
+        "data_file_base_dir": parser["data_file_base_dir"],
+        "data_list_file_path": parser["data_list_file_path"],
+        "determ": parser["determ"],
+        "learning_rate": parser["learning_rate"],
+        "learning_rate_arch": parser["learning_rate_arch"],
+        "learning_rate_milestones": np.array(parser["learning_rate_milestones"]),
+        "num_images_per_batch": parser["num_images_per_batch"],
+        "num_epochs": parser["num_epochs"],  # around 20k iterations
+        "num_epochs_per_validation": parser["num_epochs_per_validation"],
+        "num_epochs_warmup": parser["num_epochs_warmup"],
+        "num_sw_batch_size": parser["num_sw_batch_size"],
+        "output_classes": parser["output_classes"],
+        "overlap_ratio": parser["overlap_ratio"],
+        "patch_size_valid": parser["patch_size_valid"],
+        "ram_cost_factor": parser["ram_cost_factor"],
+    }
+
+    utils.init_distributed_mode(args) # Sets args.distributed among other things
+    assert args["distributed"] # don't support cycling when not distributed for simplicity
+    device = torch.device(args["device"])
+
+    # deterministic training
+    if args["determ"]:
+        set_determinism(seed=0)
+
+    timer.report('preliminaries')
+
+    train_transforms = parser.get_parsed_content("transform_train")
+    val_transforms = parser.get_parsed_content("transform_validation")
+
+    timer.report('transforms')
+
+    with open(args["data_list_file_path"], "r") as f:
+        json_data = json.load(f)
+
+    list_train = json_data["training"]
+    list_valid = json_data["validation"]
+
+    # training data
+    files = []
+    for _i in range(len(list_train)):
+        str_img = os.path.join(args["data_file_base_dir"], list_train[_i]["image"])
+        str_seg = os.path.join(args["data_file_base_dir"], list_train[_i]["label"])
+
+        if (not os.path.exists(str_img)) or (not os.path.exists(str_seg)):
+            continue
+
+        files.append({"image": str_img, "label": str_seg})
+    train_files = files
+
+    random.shuffle(train_files)
+
+    timer.report('training files')
+
+    # validation data
+    files = []
+    for _i in range(len(list_valid)):
+        str_img = os.path.join(args["data_file_base_dir"], list_valid[_i]["image"])
+        str_seg = os.path.join(args["data_file_base_dir"], list_valid[_i]["label"])
+
+        if (not os.path.exists(str_img)) or (not os.path.exists(str_seg)):
+            continue
+
+        files.append({"image": str_img, "label": str_seg})
+    val_files = files
+
+    timer.report('validation files')
+
+    n_workers = 1
+    cache_rate = 0.0
+    train_ds = monai.data.CacheDataset(
+        data=train_files, transform=train_transforms, cache_rate=cache_rate, num_workers=n_workers
+    )
+    val_ds = monai.data.CacheDataset(data=val_files, transform=val_transforms, cache_rate=cache_rate, num_workers=n_workers)
+
+    train_sampler = InterruptableDistributedSampler(train_ds)
+    val_sampler = InterruptableDistributedSampler(val_ds)
+
+    train_loader = DataLoader(train_ds, batch_size=1, sampler=train_sampler, num_workers=1)
+    val_loader = DataLoader(val_ds, batch_size=1, sampler=val_sampler, num_workers=1)
+
+    timer.report('datasets and dataloaders')
+
+    # # TESTING
+    # timer = TimestampedTimer("testing start")
+    # for i, batch_data in enumerate(train_loader):
+    #     inputs, labels = batch_data["image"], batch_data["label"]
+    #     timer.report("batch")
+    #     inputs.size == (1, 1, 96, 96, 96), labels.size == (1, 1, 96, 96, 96)
+
+    dints_space = TopologySearch(channel_mul=0.5, num_blocks=12, num_depths=4, use_downsample=True, device=device)
+    model = DiNTS(dints_space, in_channels=1, num_classes=3, use_downsample=True)
+    loss_func = DiceCELoss(include_background=False, to_onehot_y=True, softmax=True, squared_pred=True, batch=True, smooth_nr=1e-05, smooth_dr=1e-05)
+
+    model = model.to(device)
+    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+
+    post_pred = transforms.Compose([transforms.EnsureType(), transforms.AsDiscrete(to_onehot=args["output_classes"], argmax=True)])
+    post_label = transforms.Compose([transforms.EnsureType(), transforms.AsDiscrete(to_onehot=args["output_classes"])])
+
+    timer.report('model to device')
+
+    model_without_ddp = model
+    if args["distributed"]:
+        model = DistributedDataParallel(model, device_ids=[args["gpu"]], find_unused_parameters=True)
+        model_without_ddp = model.module
+
+    # optimizers
+    optimizer = torch.optim.SGD(
+        model_without_ddp.weight_parameters(), lr=args["learning_rate"] * args["world_size"], momentum=0.9, weight_decay=0.00004
+    )
+    arch_optimizer_a = torch.optim.Adam(
+        [dints_space.log_alpha_a], lr=args["learning_rate_arch"] * args["world_size"], betas=(0.5, 0.999), weight_decay=0.0
+    )
+    arch_optimizer_c = torch.optim.Adam(
+        [dints_space.log_alpha_c], lr=args["learning_rate_arch"] * args["world_size"], betas=(0.5, 0.999), weight_decay=0.0
+    )
+
+    timer.report('model ready to train')
+
+    # amp
+    if args["amp"]:
+        from torch.cuda.amp import GradScaler
+        scaler = GradScaler()
+        if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+            print("[info] amp enabled")
+
+    # start a typical PyTorch training
+    val_interval = args["num_epochs_per_validation"]
+
+    # Init metric trackers
+    train_metrics = MetricsTracker()
+    val_metric = torch.zeros((args["output_classes"] - 1) * 2, dtype=torch.float, device=device)
+
+    timer.report('metrics setup')
+
+    # RETRIEVE CHECKPOINT
+    Path(args["resume"]).parent.mkdir(parents=True, exist_ok=True)
+
+    checkpoint = None
+    if args["resume"] and os.path.isfile(args["resume"]): # If we're resuming...
+        checkpoint = torch.load(args["resume"], map_location="cpu")
+    elif args["prev_resume"] and os.path.isfile(args["prev_resume"]):
+        checkpoint = torch.load(args["prev_resume"], map_location="cpu")
+    if checkpoint is not None:
+        args["start_epoch"] = checkpoint["epoch"]
+        model_without_ddp.load_state_dict(checkpoint["model"])
+        dints_space.load_state_dict(checkpoint["dints"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        arch_optimizer_a.load_state_dict(checkpoint["arch_optimizer_a"])
+        arch_optimizer_c.load_state_dict(checkpoint["arch_optimizer_c"])
+        train_sampler.load_state_dict(checkpoint["train_sampler"])
+        val_sampler.load_state_dict(checkpoint["val_sampler"])
+        scaler.load_state_dict(checkpoint["scaler"])
+        train_metrics = checkpoint["train_metrics"]
+        val_metric = checkpoint["val_metric"]
+        val_metric.to(device)
+
+    timer.report('obtain checkpoint')
+
+    for epoch in range(args["start_epoch"], args["num_epochs"]):
+
+        print('\n')
+        print(f"EPOCH :: {epoch} / {args['num_epochs']}")
+        print('\n')
+
+        with train_sampler.in_epoch(epoch):
+            timer = TimestampedTimer("Start training")
+
+            model, dints_space, timer, train_metrics = search_one_epoch(
+                model, optimizer, dints_space, arch_optimizer_a, arch_optimizer_c, 
+                train_sampler, val_sampler, scaler, train_metrics, val_metric,
+                epoch, train_loader, loss_func, args, timer
+            )
+            timer.report(f'searching space for epoch {epoch}')
+
+            if (epoch + 1) % val_interval == 0 or (epoch + 1) == args["num_epochs"]:
+
+                with val_sampler.in_epoch(epoch):
+                    timer = TimestampedTimer("Start evaluation")
+
+                    timer = eval_search(
+                        model, optimizer, dints_space, arch_optimizer_a, arch_optimizer_c, 
+                        train_sampler, val_sampler, scaler, train_metrics, val_metric,
+                        epoch, val_loader, post_pred, post_label, args, timer
+                    )
+                    timer.report(f'evaluating search for epoch {epoch}')
+
+if __name__ == "__main__":
+    args = get_args_parser().parse_args()
+    main(args, timer)
diff --git a/monai_pancreas_dints/configs/dataset_0.json b/monai_pancreas_dints/configs/dataset_0.json
new file mode 100644
index 00000000..6f244e8d
--- /dev/null
+++ b/monai_pancreas_dints/configs/dataset_0.json
@@ -0,0 +1,1132 @@
+{
+    "training": [
+        {
+            "label": "labelsTr/pancreas_046.nii.gz",
+            "image": "imagesTr/pancreas_046.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_261.nii.gz",
+            "image": "imagesTr/pancreas_261.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_225.nii.gz",
+            "image": "imagesTr/pancreas_225.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_380.nii.gz",
+            "image": "imagesTr/pancreas_380.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_304.nii.gz",
+            "image": "imagesTr/pancreas_304.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_021.nii.gz",
+            "image": "imagesTr/pancreas_021.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_356.nii.gz",
+            "image": "imagesTr/pancreas_356.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_387.nii.gz",
+            "image": "imagesTr/pancreas_387.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_050.nii.gz",
+            "image": "imagesTr/pancreas_050.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_024.nii.gz",
+            "image": "imagesTr/pancreas_024.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_242.nii.gz",
+            "image": "imagesTr/pancreas_242.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_103.nii.gz",
+            "image": "imagesTr/pancreas_103.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_130.nii.gz",
+            "image": "imagesTr/pancreas_130.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_404.nii.gz",
+            "image": "imagesTr/pancreas_404.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_094.nii.gz",
+            "image": "imagesTr/pancreas_094.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_360.nii.gz",
+            "image": "imagesTr/pancreas_360.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_155.nii.gz",
+            "image": "imagesTr/pancreas_155.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_181.nii.gz",
+            "image": "imagesTr/pancreas_181.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_067.nii.gz",
+            "image": "imagesTr/pancreas_067.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_413.nii.gz",
+            "image": "imagesTr/pancreas_413.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_256.nii.gz",
+            "image": "imagesTr/pancreas_256.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_355.nii.gz",
+            "image": "imagesTr/pancreas_355.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_310.nii.gz",
+            "image": "imagesTr/pancreas_310.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_392.nii.gz",
+            "image": "imagesTr/pancreas_392.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_320.nii.gz",
+            "image": "imagesTr/pancreas_320.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_089.nii.gz",
+            "image": "imagesTr/pancreas_089.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_370.nii.gz",
+            "image": "imagesTr/pancreas_370.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_358.nii.gz",
+            "image": "imagesTr/pancreas_358.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_201.nii.gz",
+            "image": "imagesTr/pancreas_201.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_278.nii.gz",
+            "image": "imagesTr/pancreas_278.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_235.nii.gz",
+            "image": "imagesTr/pancreas_235.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_418.nii.gz",
+            "image": "imagesTr/pancreas_418.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_318.nii.gz",
+            "image": "imagesTr/pancreas_318.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_374.nii.gz",
+            "image": "imagesTr/pancreas_374.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_234.nii.gz",
+            "image": "imagesTr/pancreas_234.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_343.nii.gz",
+            "image": "imagesTr/pancreas_343.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_302.nii.gz",
+            "image": "imagesTr/pancreas_302.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_126.nii.gz",
+            "image": "imagesTr/pancreas_126.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_077.nii.gz",
+            "image": "imagesTr/pancreas_077.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_361.nii.gz",
+            "image": "imagesTr/pancreas_361.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_326.nii.gz",
+            "image": "imagesTr/pancreas_326.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_401.nii.gz",
+            "image": "imagesTr/pancreas_401.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_241.nii.gz",
+            "image": "imagesTr/pancreas_241.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_028.nii.gz",
+            "image": "imagesTr/pancreas_028.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_414.nii.gz",
+            "image": "imagesTr/pancreas_414.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_364.nii.gz",
+            "image": "imagesTr/pancreas_364.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_292.nii.gz",
+            "image": "imagesTr/pancreas_292.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_259.nii.gz",
+            "image": "imagesTr/pancreas_259.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_244.nii.gz",
+            "image": "imagesTr/pancreas_244.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_018.nii.gz",
+            "image": "imagesTr/pancreas_018.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_125.nii.gz",
+            "image": "imagesTr/pancreas_125.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_419.nii.gz",
+            "image": "imagesTr/pancreas_419.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_296.nii.gz",
+            "image": "imagesTr/pancreas_296.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_315.nii.gz",
+            "image": "imagesTr/pancreas_315.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_200.nii.gz",
+            "image": "imagesTr/pancreas_200.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_389.nii.gz",
+            "image": "imagesTr/pancreas_389.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_104.nii.gz",
+            "image": "imagesTr/pancreas_104.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_182.nii.gz",
+            "image": "imagesTr/pancreas_182.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_230.nii.gz",
+            "image": "imagesTr/pancreas_230.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_147.nii.gz",
+            "image": "imagesTr/pancreas_147.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_344.nii.gz",
+            "image": "imagesTr/pancreas_344.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_308.nii.gz",
+            "image": "imagesTr/pancreas_308.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_081.nii.gz",
+            "image": "imagesTr/pancreas_081.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_298.nii.gz",
+            "image": "imagesTr/pancreas_298.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_275.nii.gz",
+            "image": "imagesTr/pancreas_275.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_106.nii.gz",
+            "image": "imagesTr/pancreas_106.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_101.nii.gz",
+            "image": "imagesTr/pancreas_101.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_158.nii.gz",
+            "image": "imagesTr/pancreas_158.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_382.nii.gz",
+            "image": "imagesTr/pancreas_382.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_393.nii.gz",
+            "image": "imagesTr/pancreas_393.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_325.nii.gz",
+            "image": "imagesTr/pancreas_325.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_001.nii.gz",
+            "image": "imagesTr/pancreas_001.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_083.nii.gz",
+            "image": "imagesTr/pancreas_083.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_249.nii.gz",
+            "image": "imagesTr/pancreas_249.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_075.nii.gz",
+            "image": "imagesTr/pancreas_075.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_100.nii.gz",
+            "image": "imagesTr/pancreas_100.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_365.nii.gz",
+            "image": "imagesTr/pancreas_365.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_084.nii.gz",
+            "image": "imagesTr/pancreas_084.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_300.nii.gz",
+            "image": "imagesTr/pancreas_300.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_339.nii.gz",
+            "image": "imagesTr/pancreas_339.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_178.nii.gz",
+            "image": "imagesTr/pancreas_178.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_219.nii.gz",
+            "image": "imagesTr/pancreas_219.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_035.nii.gz",
+            "image": "imagesTr/pancreas_035.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_388.nii.gz",
+            "image": "imagesTr/pancreas_388.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_276.nii.gz",
+            "image": "imagesTr/pancreas_276.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_197.nii.gz",
+            "image": "imagesTr/pancreas_197.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_379.nii.gz",
+            "image": "imagesTr/pancreas_379.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_120.nii.gz",
+            "image": "imagesTr/pancreas_120.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_186.nii.gz",
+            "image": "imagesTr/pancreas_186.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_236.nii.gz",
+            "image": "imagesTr/pancreas_236.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_071.nii.gz",
+            "image": "imagesTr/pancreas_071.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_051.nii.gz",
+            "image": "imagesTr/pancreas_051.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_229.nii.gz",
+            "image": "imagesTr/pancreas_229.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_203.nii.gz",
+            "image": "imagesTr/pancreas_203.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_330.nii.gz",
+            "image": "imagesTr/pancreas_330.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_145.nii.gz",
+            "image": "imagesTr/pancreas_145.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_019.nii.gz",
+            "image": "imagesTr/pancreas_019.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_015.nii.gz",
+            "image": "imagesTr/pancreas_015.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_293.nii.gz",
+            "image": "imagesTr/pancreas_293.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_386.nii.gz",
+            "image": "imagesTr/pancreas_386.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_299.nii.gz",
+            "image": "imagesTr/pancreas_299.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_421.nii.gz",
+            "image": "imagesTr/pancreas_421.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_412.nii.gz",
+            "image": "imagesTr/pancreas_412.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_243.nii.gz",
+            "image": "imagesTr/pancreas_243.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_228.nii.gz",
+            "image": "imagesTr/pancreas_228.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_109.nii.gz",
+            "image": "imagesTr/pancreas_109.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_006.nii.gz",
+            "image": "imagesTr/pancreas_006.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_140.nii.gz",
+            "image": "imagesTr/pancreas_140.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_284.nii.gz",
+            "image": "imagesTr/pancreas_284.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_372.nii.gz",
+            "image": "imagesTr/pancreas_372.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_087.nii.gz",
+            "image": "imagesTr/pancreas_087.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_198.nii.gz",
+            "image": "imagesTr/pancreas_198.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_277.nii.gz",
+            "image": "imagesTr/pancreas_277.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_091.nii.gz",
+            "image": "imagesTr/pancreas_091.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_004.nii.gz",
+            "image": "imagesTr/pancreas_004.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_040.nii.gz",
+            "image": "imagesTr/pancreas_040.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_025.nii.gz",
+            "image": "imagesTr/pancreas_025.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_088.nii.gz",
+            "image": "imagesTr/pancreas_088.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_096.nii.gz",
+            "image": "imagesTr/pancreas_096.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_285.nii.gz",
+            "image": "imagesTr/pancreas_285.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_049.nii.gz",
+            "image": "imagesTr/pancreas_049.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_291.nii.gz",
+            "image": "imagesTr/pancreas_291.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_172.nii.gz",
+            "image": "imagesTr/pancreas_172.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_294.nii.gz",
+            "image": "imagesTr/pancreas_294.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_080.nii.gz",
+            "image": "imagesTr/pancreas_080.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_214.nii.gz",
+            "image": "imagesTr/pancreas_214.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_309.nii.gz",
+            "image": "imagesTr/pancreas_309.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_290.nii.gz",
+            "image": "imagesTr/pancreas_290.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_313.nii.gz",
+            "image": "imagesTr/pancreas_313.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_175.nii.gz",
+            "image": "imagesTr/pancreas_175.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_029.nii.gz",
+            "image": "imagesTr/pancreas_029.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_222.nii.gz",
+            "image": "imagesTr/pancreas_222.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_395.nii.gz",
+            "image": "imagesTr/pancreas_395.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_253.nii.gz",
+            "image": "imagesTr/pancreas_253.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_196.nii.gz",
+            "image": "imagesTr/pancreas_196.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_183.nii.gz",
+            "image": "imagesTr/pancreas_183.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_113.nii.gz",
+            "image": "imagesTr/pancreas_113.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_376.nii.gz",
+            "image": "imagesTr/pancreas_376.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_258.nii.gz",
+            "image": "imagesTr/pancreas_258.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_367.nii.gz",
+            "image": "imagesTr/pancreas_367.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_032.nii.gz",
+            "image": "imagesTr/pancreas_032.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_239.nii.gz",
+            "image": "imagesTr/pancreas_239.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_316.nii.gz",
+            "image": "imagesTr/pancreas_316.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_045.nii.gz",
+            "image": "imagesTr/pancreas_045.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_227.nii.gz",
+            "image": "imagesTr/pancreas_227.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_099.nii.gz",
+            "image": "imagesTr/pancreas_099.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_264.nii.gz",
+            "image": "imagesTr/pancreas_264.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_122.nii.gz",
+            "image": "imagesTr/pancreas_122.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_105.nii.gz",
+            "image": "imagesTr/pancreas_105.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_111.nii.gz",
+            "image": "imagesTr/pancreas_111.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_415.nii.gz",
+            "image": "imagesTr/pancreas_415.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_217.nii.gz",
+            "image": "imagesTr/pancreas_217.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_270.nii.gz",
+            "image": "imagesTr/pancreas_270.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_124.nii.gz",
+            "image": "imagesTr/pancreas_124.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_411.nii.gz",
+            "image": "imagesTr/pancreas_411.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_066.nii.gz",
+            "image": "imagesTr/pancreas_066.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_254.nii.gz",
+            "image": "imagesTr/pancreas_254.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_159.nii.gz",
+            "image": "imagesTr/pancreas_159.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_199.nii.gz",
+            "image": "imagesTr/pancreas_199.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_328.nii.gz",
+            "image": "imagesTr/pancreas_328.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_166.nii.gz",
+            "image": "imagesTr/pancreas_166.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_173.nii.gz",
+            "image": "imagesTr/pancreas_173.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_287.nii.gz",
+            "image": "imagesTr/pancreas_287.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_157.nii.gz",
+            "image": "imagesTr/pancreas_157.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_209.nii.gz",
+            "image": "imagesTr/pancreas_209.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_169.nii.gz",
+            "image": "imagesTr/pancreas_169.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_179.nii.gz",
+            "image": "imagesTr/pancreas_179.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_224.nii.gz",
+            "image": "imagesTr/pancreas_224.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_283.nii.gz",
+            "image": "imagesTr/pancreas_283.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_297.nii.gz",
+            "image": "imagesTr/pancreas_297.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_377.nii.gz",
+            "image": "imagesTr/pancreas_377.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_405.nii.gz",
+            "image": "imagesTr/pancreas_405.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_137.nii.gz",
+            "image": "imagesTr/pancreas_137.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_074.nii.gz",
+            "image": "imagesTr/pancreas_074.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_005.nii.gz",
+            "image": "imagesTr/pancreas_005.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_110.nii.gz",
+            "image": "imagesTr/pancreas_110.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_321.nii.gz",
+            "image": "imagesTr/pancreas_321.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_410.nii.gz",
+            "image": "imagesTr/pancreas_410.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_385.nii.gz",
+            "image": "imagesTr/pancreas_385.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_187.nii.gz",
+            "image": "imagesTr/pancreas_187.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_345.nii.gz",
+            "image": "imagesTr/pancreas_345.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_061.nii.gz",
+            "image": "imagesTr/pancreas_061.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_119.nii.gz",
+            "image": "imagesTr/pancreas_119.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_086.nii.gz",
+            "image": "imagesTr/pancreas_086.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_211.nii.gz",
+            "image": "imagesTr/pancreas_211.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_193.nii.gz",
+            "image": "imagesTr/pancreas_193.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_165.nii.gz",
+            "image": "imagesTr/pancreas_165.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_346.nii.gz",
+            "image": "imagesTr/pancreas_346.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_329.nii.gz",
+            "image": "imagesTr/pancreas_329.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_098.nii.gz",
+            "image": "imagesTr/pancreas_098.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_207.nii.gz",
+            "image": "imagesTr/pancreas_207.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_180.nii.gz",
+            "image": "imagesTr/pancreas_180.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_135.nii.gz",
+            "image": "imagesTr/pancreas_135.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_037.nii.gz",
+            "image": "imagesTr/pancreas_037.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_354.nii.gz",
+            "image": "imagesTr/pancreas_354.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_167.nii.gz",
+            "image": "imagesTr/pancreas_167.nii.gz"
+        }
+    ],
+    "validation": [
+        {
+            "label": "labelsTr/pancreas_398.nii.gz",
+            "image": "imagesTr/pancreas_398.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_336.nii.gz",
+            "image": "imagesTr/pancreas_336.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_347.nii.gz",
+            "image": "imagesTr/pancreas_347.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_416.nii.gz",
+            "image": "imagesTr/pancreas_416.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_268.nii.gz",
+            "image": "imagesTr/pancreas_268.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_204.nii.gz",
+            "image": "imagesTr/pancreas_204.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_070.nii.gz",
+            "image": "imagesTr/pancreas_070.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_327.nii.gz",
+            "image": "imagesTr/pancreas_327.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_210.nii.gz",
+            "image": "imagesTr/pancreas_210.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_333.nii.gz",
+            "image": "imagesTr/pancreas_333.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_274.nii.gz",
+            "image": "imagesTr/pancreas_274.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_191.nii.gz",
+            "image": "imagesTr/pancreas_191.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_058.nii.gz",
+            "image": "imagesTr/pancreas_058.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_117.nii.gz",
+            "image": "imagesTr/pancreas_117.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_016.nii.gz",
+            "image": "imagesTr/pancreas_016.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_218.nii.gz",
+            "image": "imagesTr/pancreas_218.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_215.nii.gz",
+            "image": "imagesTr/pancreas_215.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_267.nii.gz",
+            "image": "imagesTr/pancreas_267.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_311.nii.gz",
+            "image": "imagesTr/pancreas_311.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_357.nii.gz",
+            "image": "imagesTr/pancreas_357.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_266.nii.gz",
+            "image": "imagesTr/pancreas_266.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_041.nii.gz",
+            "image": "imagesTr/pancreas_041.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_255.nii.gz",
+            "image": "imagesTr/pancreas_255.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_043.nii.gz",
+            "image": "imagesTr/pancreas_043.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_102.nii.gz",
+            "image": "imagesTr/pancreas_102.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_280.nii.gz",
+            "image": "imagesTr/pancreas_280.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_247.nii.gz",
+            "image": "imagesTr/pancreas_247.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_362.nii.gz",
+            "image": "imagesTr/pancreas_362.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_265.nii.gz",
+            "image": "imagesTr/pancreas_265.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_138.nii.gz",
+            "image": "imagesTr/pancreas_138.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_399.nii.gz",
+            "image": "imagesTr/pancreas_399.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_391.nii.gz",
+            "image": "imagesTr/pancreas_391.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_064.nii.gz",
+            "image": "imagesTr/pancreas_064.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_010.nii.gz",
+            "image": "imagesTr/pancreas_010.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_127.nii.gz",
+            "image": "imagesTr/pancreas_127.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_092.nii.gz",
+            "image": "imagesTr/pancreas_092.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_246.nii.gz",
+            "image": "imagesTr/pancreas_246.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_213.nii.gz",
+            "image": "imagesTr/pancreas_213.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_148.nii.gz",
+            "image": "imagesTr/pancreas_148.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_160.nii.gz",
+            "image": "imagesTr/pancreas_160.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_042.nii.gz",
+            "image": "imagesTr/pancreas_042.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_409.nii.gz",
+            "image": "imagesTr/pancreas_409.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_342.nii.gz",
+            "image": "imagesTr/pancreas_342.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_107.nii.gz",
+            "image": "imagesTr/pancreas_107.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_348.nii.gz",
+            "image": "imagesTr/pancreas_348.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_406.nii.gz",
+            "image": "imagesTr/pancreas_406.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_279.nii.gz",
+            "image": "imagesTr/pancreas_279.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_301.nii.gz",
+            "image": "imagesTr/pancreas_301.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_378.nii.gz",
+            "image": "imagesTr/pancreas_378.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_295.nii.gz",
+            "image": "imagesTr/pancreas_295.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_056.nii.gz",
+            "image": "imagesTr/pancreas_056.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_226.nii.gz",
+            "image": "imagesTr/pancreas_226.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_048.nii.gz",
+            "image": "imagesTr/pancreas_048.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_375.nii.gz",
+            "image": "imagesTr/pancreas_375.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_366.nii.gz",
+            "image": "imagesTr/pancreas_366.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_334.nii.gz",
+            "image": "imagesTr/pancreas_334.nii.gz"
+        }
+    ],
+    "testing": [
+        {
+            "label": "labelsTr/pancreas_212.nii.gz",
+            "image": "imagesTr/pancreas_212.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_131.nii.gz",
+            "image": "imagesTr/pancreas_131.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_400.nii.gz",
+            "image": "imagesTr/pancreas_400.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_170.nii.gz",
+            "image": "imagesTr/pancreas_170.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_312.nii.gz",
+            "image": "imagesTr/pancreas_312.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_069.nii.gz",
+            "image": "imagesTr/pancreas_069.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_351.nii.gz",
+            "image": "imagesTr/pancreas_351.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_095.nii.gz",
+            "image": "imagesTr/pancreas_095.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_231.nii.gz",
+            "image": "imagesTr/pancreas_231.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_303.nii.gz",
+            "image": "imagesTr/pancreas_303.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_305.nii.gz",
+            "image": "imagesTr/pancreas_305.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_055.nii.gz",
+            "image": "imagesTr/pancreas_055.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_350.nii.gz",
+            "image": "imagesTr/pancreas_350.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_286.nii.gz",
+            "image": "imagesTr/pancreas_286.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_269.nii.gz",
+            "image": "imagesTr/pancreas_269.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_331.nii.gz",
+            "image": "imagesTr/pancreas_331.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_093.nii.gz",
+            "image": "imagesTr/pancreas_093.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_114.nii.gz",
+            "image": "imagesTr/pancreas_114.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_402.nii.gz",
+            "image": "imagesTr/pancreas_402.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_369.nii.gz",
+            "image": "imagesTr/pancreas_369.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_078.nii.gz",
+            "image": "imagesTr/pancreas_078.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_262.nii.gz",
+            "image": "imagesTr/pancreas_262.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_052.nii.gz",
+            "image": "imagesTr/pancreas_052.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_129.nii.gz",
+            "image": "imagesTr/pancreas_129.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_194.nii.gz",
+            "image": "imagesTr/pancreas_194.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_289.nii.gz",
+            "image": "imagesTr/pancreas_289.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_012.nii.gz",
+            "image": "imagesTr/pancreas_012.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_149.nii.gz",
+            "image": "imagesTr/pancreas_149.nii.gz"
+        },
+        {
+            "label": "labelsTr/pancreas_323.nii.gz",
+            "image": "imagesTr/pancreas_323.nii.gz"
+        }
+    ]
+}
\ No newline at end of file
diff --git a/monai_pancreas_dints/configs/evaluate.yaml b/monai_pancreas_dints/configs/evaluate.yaml
new file mode 100644
index 00000000..520c57e7
--- /dev/null
+++ b/monai_pancreas_dints/configs/evaluate.yaml
@@ -0,0 +1,65 @@
+---
+validate#postprocessing:
+  _target_: Compose
+  transforms:
+  - _target_: Activationsd
+    keys: pred
+    softmax: true
+  - _target_: Invertd
+    keys:
+    - pred
+    - label
+    transform: "@validate#preprocessing"
+    orig_keys: image
+    meta_key_postfix: meta_dict
+    nearest_interp:
+    - false
+    - true
+    to_tensor: true
+  - _target_: AsDiscreted
+    keys:
+    - pred
+    - label
+    argmax:
+    - true
+    - false
+    to_onehot: 3
+  - _target_: CopyItemsd
+    keys: "pred"
+    times: 1
+    names: "pred_save"
+  - _target_: AsDiscreted
+    keys:
+    - pred_save
+    argmax:
+    - true
+  - _target_: SaveImaged
+    keys: pred_save
+    meta_keys: pred_meta_dict
+    output_dir: "@output_dir"
+    resample: false
+    squeeze_end_dims: true
+validate#dataset:
+    _target_: Dataset
+    data: "@val_datalist"
+    transform: "@validate#preprocessing"
+validate#handlers:
+- _target_: CheckpointLoader
+  load_path: "$@ckpt_dir + '/model.pt'"
+  load_dict:
+    model: "@network"
+- _target_: StatsHandler
+  iteration_log: false
+- _target_: MetricsSaver
+  save_dir: "@output_dir"
+  metrics:
+  - val_mean_dice
+  - val_acc
+  metric_details:
+  - val_mean_dice
+  batch_transform: "$monai.handlers.from_engine(['image_meta_dict'])"
+  summary_ops: "*"
+initialize:
+- "$setattr(torch.backends.cudnn, 'benchmark', True)"
+run:
+- "$@validate#evaluator.run()"
diff --git a/monai_pancreas_dints/configs/inference.yaml b/monai_pancreas_dints/configs/inference.yaml
new file mode 100644
index 00000000..d48732b8
--- /dev/null
+++ b/monai_pancreas_dints/configs/inference.yaml
@@ -0,0 +1,117 @@
+---
+imports:
+- "$import glob"
+- "$import os"
+input_channels: 1
+output_classes: 3
+arch_ckpt_path: "$@bundle_root + '/models/search_code_18590.pt'"
+arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=torch.device('cuda'))"
+bundle_root: "."
+output_dir: "$@bundle_root + '/eval'"
+dataset_dir: "/workspace/data/msd/Task07_Pancreas"
+data_list_file_path: "$@bundle_root + '/configs/dataset_0.json'"
+datalist: "$monai.data.load_decathlon_datalist(@data_list_file_path, data_list_key='testing',
+  base_dir=@dataset_dir)"
+device: "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
+dints_space:
+  _target_: monai.networks.nets.TopologyInstance
+  channel_mul: 1
+  num_blocks: 12
+  num_depths: 4
+  use_downsample: true
+  arch_code:
+  - "$@arch_ckpt['arch_code_a']"
+  - "$@arch_ckpt['arch_code_c']"
+  device: "$torch.device('cuda')"
+network_def:
+  _target_: monai.networks.nets.DiNTS
+  dints_space: "@dints_space"
+  in_channels: "@input_channels"
+  num_classes: "@output_classes"
+  use_downsample: true
+  node_a: "$torch.from_numpy(@arch_ckpt['node_a'])"
+network: "$@network_def.to(@device)"
+preprocessing:
+  _target_: Compose
+  transforms:
+  - _target_: LoadImaged
+    keys: image
+  - _target_: EnsureChannelFirstd
+    keys: image
+  - _target_: Orientationd
+    keys: image
+    axcodes: RAS
+  - _target_: Spacingd
+    keys: image
+    pixdim:
+    - 1
+    - 1
+    - 1
+    mode: bilinear
+  - _target_: ScaleIntensityRanged
+    keys: image
+    a_min: -87
+    a_max: 199
+    b_min: 0
+    b_max: 1
+    clip: true
+  - _target_: EnsureTyped
+    keys: image
+dataset:
+  _target_: Dataset
+  data: "@datalist"
+  transform: "@preprocessing"
+dataloader:
+  _target_: DataLoader
+  dataset: "@dataset"
+  batch_size: 1
+  shuffle: false
+  num_workers: 4
+inferer:
+  _target_: SlidingWindowInferer
+  roi_size:
+  - 96
+  - 96
+  - 96
+  sw_batch_size: 4
+  overlap: 0.625
+postprocessing:
+  _target_: Compose
+  transforms:
+  - _target_: Activationsd
+    keys: pred
+    softmax: true
+  - _target_: Invertd
+    keys: pred
+    transform: "@preprocessing"
+    orig_keys: image
+    meta_key_postfix: meta_dict
+    nearest_interp: false
+    to_tensor: true
+  - _target_: AsDiscreted
+    keys: pred
+    argmax: true
+  - _target_: SaveImaged
+    keys: pred
+    meta_keys: pred_meta_dict
+    output_dir: "@output_dir"
+handlers:
+- _target_: CheckpointLoader
+  load_path: "$@bundle_root + '/models/model.pt'"
+  load_dict:
+    model: "@network"
+- _target_: StatsHandler
+  iteration_log: false
+evaluator:
+  _target_: SupervisedEvaluator
+  device: "@device"
+  val_data_loader: "@dataloader"
+  network: "@network"
+  inferer: "@inferer"
+  postprocessing: "@postprocessing"
+  val_handlers: "@handlers"
+  amp: true
+initialize:
+- "$setattr(torch.backends.cudnn, 'benchmark', True)"
+run:
+- "$@evaluator.run()"
diff --git a/monai_pancreas_dints/configs/inference_trt.yaml b/monai_pancreas_dints/configs/inference_trt.yaml
new file mode 100644
index 00000000..1bb4820d
--- /dev/null
+++ b/monai_pancreas_dints/configs/inference_trt.yaml
@@ -0,0 +1,8 @@
+---
+imports:
+- "$import glob"
+- "$import os"
+- "$import torch_tensorrt"
+handlers#0#_disabled_: true
+network_def: "$torch.jit.load(@bundle_root + '/models/model_trt.ts')"
+evaluator#amp: false
diff --git a/monai_pancreas_dints/configs/logging.conf b/monai_pancreas_dints/configs/logging.conf
new file mode 100644
index 00000000..91c1a21c
--- /dev/null
+++ b/monai_pancreas_dints/configs/logging.conf
@@ -0,0 +1,21 @@
+[loggers]
+keys=root
+
+[handlers]
+keys=consoleHandler
+
+[formatters]
+keys=fullFormatter
+
+[logger_root]
+level=INFO
+handlers=consoleHandler
+
+[handler_consoleHandler]
+class=StreamHandler
+level=INFO
+formatter=fullFormatter
+args=(sys.stdout,)
+
+[formatter_fullFormatter]
+format=%(asctime)s - %(name)s - %(levelname)s - %(message)s
diff --git a/monai_pancreas_dints/configs/metadata.json b/monai_pancreas_dints/configs/metadata.json
new file mode 100644
index 00000000..2fd6955f
--- /dev/null
+++ b/monai_pancreas_dints/configs/metadata.json
@@ -0,0 +1,97 @@
+{
+    "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20220324.json",
+    "version": "0.4.3",
+    "changelog": {
+        "0.4.3": "add support for TensorRT conversion and inference",
+        "0.4.2": "update search function to match monai 1.2",
+        "0.4.1": "fix the wrong GPU index issue of multi-node",
+        "0.4.0": "remove error dollar symbol in readme",
+        "0.3.9": "add cpu ram requirement in readme",
+        "0.3.8": "add non-deterministic note",
+        "0.3.7": "re-train model with updated dints implementation",
+        "0.3.6": "black autofix format and add name tag",
+        "0.3.5": "restructure readme to match updated template",
+        "0.3.4": "correct typos",
+        "0.3.3": "update learning rate and readme",
+        "0.3.2": "update to use monai 1.0.1",
+        "0.3.1": "fix license Copyright error",
+        "0.3.0": "update license files",
+        "0.2.0": "unify naming",
+        "0.1.1": "fix data type issue in searching/training configurations",
+        "0.1.0": "complete the model package",
+        "0.0.1": "initialize the model package structure"
+    },
+    "monai_version": "1.2.0",
+    "pytorch_version": "1.13.1",
+    "numpy_version": "1.22.2",
+    "optional_packages_version": {
+        "fire": "0.4.0",
+        "nibabel": "4.0.1",
+        "pytorch-ignite": "0.4.9"
+    },
+    "name": "Pancreas CT DiNTS segmentation",
+    "task": "Neural architecture search on pancreas CT segmentation",
+    "description": "Searched architectures for volumetric (3D) segmentation of the pancreas from CT image",
+    "authors": "MONAI team",
+    "copyright": "Copyright (c) MONAI Consortium",
+    "data_source": "Task07_Pancreas.tar from http://medicaldecathlon.com/",
+    "data_type": "nibabel",
+    "image_classes": "single channel data, intensity scaled to [0, 1]",
+    "label_classes": "single channel data, 1 is pancreas, 2 is pancreatic tumor, 0 is everything else",
+    "pred_classes": "3 channels OneHot data, channel 1 is pancreas, channel 2 is pancreatic tumor, channel 0 is background",
+    "eval_metrics": {
+        "mean_dice": 0.62
+    },
+    "intended_use": "This is an example, not to be used for diagnostic purposes",
+    "references": [
+        "He, Y., Yang, D., Roth, H., Zhao, C. and Xu, D., 2021. Dints: Differentiable neural network topology search for 3d medical image segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 5841-5850)."
+    ],
+    "network_data_format": {
+        "inputs": {
+            "image": {
+                "type": "image",
+                "format": "hounsfield",
+                "modality": "CT",
+                "num_channels": 1,
+                "spatial_shape": [
+                    96,
+                    96,
+                    96
+                ],
+                "dtype": "float32",
+                "value_range": [
+                    0,
+                    1
+                ],
+                "is_patch_data": true,
+                "channel_def": {
+                    "0": "image"
+                }
+            }
+        },
+        "outputs": {
+            "pred": {
+                "type": "image",
+                "format": "segmentation",
+                "num_channels": 3,
+                "spatial_shape": [
+                    96,
+                    96,
+                    96
+                ],
+                "dtype": "float32",
+                "value_range": [
+                    0,
+                    1,
+                    2
+                ],
+                "is_patch_data": true,
+                "channel_def": {
+                    "0": "background",
+                    "1": "pancreas",
+                    "2": "pancreatic tumor"
+                }
+            }
+        }
+    }
+}
diff --git a/monai_pancreas_dints/configs/multi_gpu_train.yaml b/monai_pancreas_dints/configs/multi_gpu_train.yaml
new file mode 100644
index 00000000..bce693fe
--- /dev/null
+++ b/monai_pancreas_dints/configs/multi_gpu_train.yaml
@@ -0,0 +1,54 @@
+---
+device: "$torch.device('cuda:' + os.environ['LOCAL_RANK'])"
+network:
+  _target_: torch.nn.parallel.DistributedDataParallel
+  module: "$@network_def.to(@device)"
+  find_unused_parameters: true
+  device_ids:
+  - "@device"
+optimizer#lr: "$0.025*dist.get_world_size()"
+lr_scheduler#step_size: "$80*dist.get_world_size()"
+train#handlers:
+  - _target_: LrScheduleHandler
+    lr_scheduler: "@lr_scheduler"
+    print_lr: true
+  - _target_: ValidationHandler
+    validator: "@validate#evaluator"
+    epoch_level: true
+    interval: "$10*dist.get_world_size()"
+  - _target_: StatsHandler
+    tag_name: train_loss
+    output_transform: "$monai.handlers.from_engine(['loss'], first=True)"
+  - _target_: TensorBoardStatsHandler
+    log_dir: "@output_dir"
+    tag_name: train_loss
+    output_transform: "$monai.handlers.from_engine(['loss'], first=True)"
+train#trainer#max_epochs: "$400*dist.get_world_size()"
+train#trainer#train_handlers: "$@train#handlers[: -2 if dist.get_rank() > 0 else None]"
+validate#evaluator#val_handlers: "$None if dist.get_rank() > 0 else @validate#handlers"
+initialize:
+- "$import torch.distributed as dist"
+- "$dist.is_initialized() or dist.init_process_group(backend='nccl')"
+- "$torch.cuda.set_device(@device)"
+- "$monai.utils.set_determinism(seed=123)"
+- "$setattr(torch.backends.cudnn, 'benchmark', True)"
+run:
+- "$@train#trainer.run()"
+finalize:
+- "$dist.is_initialized() and dist.destroy_process_group()"
+train_data_partition: "$monai.data.partition_dataset(data=@train_datalist, num_partitions=dist.get_world_size(),
+  shuffle=True, even_divisible=True,)[dist.get_rank()]"
+train#dataset:
+  _target_: CacheDataset
+  data: "@train_data_partition"
+  transform: "@train#preprocessing"
+  cache_rate: 1
+  num_workers: 4
+val_data_partition: "$monai.data.partition_dataset(data=@val_datalist, num_partitions=dist.get_world_size(),
+  shuffle=False, even_divisible=False,)[dist.get_rank()]"
+validate#dataset:
+  _target_: CacheDataset
+  data: "@val_data_partition"
+  transform: "@validate#preprocessing"
+  cache_rate: 1
+  num_workers: 4
diff --git a/monai_pancreas_dints/configs/search.yaml b/monai_pancreas_dints/configs/search.yaml
new file mode 100644
index 00000000..52525500
--- /dev/null
+++ b/monai_pancreas_dints/configs/search.yaml
@@ -0,0 +1,278 @@
+---
+imports:
+  - "$from scipy import ndimage"
+arch_ckpt_path: models
+amp: true
+# data_file_base_dir: /workspace/data/msd/Task07_Pancreas
+data_file_base_dir: /mnt/Datasets/Open-Datasets/MONAI/Task07_Pancreas
+data_list_file_path: configs/dataset_0.json
+determ: true
+input_channels: 1
+learning_rate: 0.0025
+learning_rate_arch: 0.0001
+learning_rate_milestones:
+- 0.4
+- 0.8
+num_images_per_batch: 1
+num_epochs: 1430
+num_epochs_per_validation: 1
+num_epochs_warmup: 715
+num_patches_per_image: 1
+num_sw_batch_size: 6
+output_classes: 3
+overlap_ratio: 0.625
+patch_size:
+- 96
+- 96
+- 96
+patch_size_valid:
+- 96
+- 96
+- 96
+ram_cost_factor: 0.8
+image_key: image
+label_key: label
+transform_train:
+  _target_: Compose
+  transforms:
+  - _target_: LoadImaged
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: EnsureChannelFirstd
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: Orientationd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    axcodes: RAS
+  - _target_: Spacingd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    pixdim:
+    - 1
+    - 1
+    - 1
+    mode:
+    - bilinear
+    - nearest
+    align_corners:
+    - true
+    - true
+  - _target_: CastToTyped
+    keys: "@image_key"
+    dtype: "$torch.float32"
+  - _target_: ScaleIntensityRanged
+    keys: "@image_key"
+    a_min: -87
+    a_max: 199
+    b_min: 0
+    b_max: 1
+    clip: true
+  - _target_: CastToTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+    dtype:
+    - "$np.float16"
+    - "$np.uint8"
+  - _target_: CopyItemsd
+    keys: "@label_key"
+    times: 1
+    names:
+    - label4crop
+  - _target_: Lambdad
+    keys: label4crop
+    func: "$lambda x, s=@output_classes: np.concatenate(tuple([ndimage.binary_dilation((x==_k).astype(x.dtype), iterations=48).astype(float) for _k in range(s)]), axis=0)"
+    overwrite: true
+  - _target_: EnsureTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: CastToTyped
+    keys: "@image_key"
+    dtype: "$torch.float32"
+  - _target_: SpatialPadd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    - label4crop
+    spatial_size: "@patch_size"
+    mode:
+    - reflect
+    - constant
+    - constant
+  - _target_: RandCropByLabelClassesd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    label_key: label4crop
+    num_classes: "@output_classes"
+    ratios: "$[1,] * @output_classes"
+    spatial_size: "@patch_size"
+    num_samples: "@num_patches_per_image"
+  - _target_: Lambdad
+    keys: label4crop
+    func: "$lambda x: 0"
+  - _target_: RandRotated
+    keys:
+    - "@image_key"
+    - "@label_key"
+    range_x: 0.3
+    range_y: 0.3
+    range_z: 0.3
+    mode:
+    - bilinear
+    - nearest
+    prob: 0.2
+  - _target_: RandZoomd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    min_zoom: 0.8
+    max_zoom: 1.2
+    mode:
+    - trilinear
+    - nearest
+    align_corners:
+    - null
+    - null
+    prob: 0.16
+  - _target_: RandGaussianSmoothd
+    keys: "@image_key"
+    sigma_x:
+    - 0.5
+    - 1.15
+    sigma_y:
+    - 0.5
+    - 1.15
+    sigma_z:
+    - 0.5
+    - 1.15
+    prob: 0.15
+  - _target_: RandScaleIntensityd
+    keys: "@image_key"
+    factors: 0.3
+    prob: 0.5
+  - _target_: RandShiftIntensityd
+    keys: "@image_key"
+    offsets: 0.1
+    prob: 0.5
+  - _target_: RandGaussianNoised
+    keys: "@image_key"
+    std: 0.01
+    prob: 0.15
+  - _target_: RandFlipd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    spatial_axis: 0
+    prob: 0.5
+  - _target_: RandFlipd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    spatial_axis: 1
+    prob: 0.5
+  - _target_: RandFlipd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    spatial_axis: 2
+    prob: 0.5
+  - _target_: CastToTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+    dtype:
+    - "$torch.float32"
+    - "$torch.uint8"
+  - _target_: ToTensord
+    keys:
+    - "@image_key"
+    - "@label_key"
+transform_validation:
+  _target_: Compose
+  transforms:
+  - _target_: LoadImaged
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: EnsureChannelFirstd
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: Orientationd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    axcodes: RAS
+  - _target_: Spacingd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    pixdim:
+    - 1
+    - 1
+    - 1
+    mode:
+    - bilinear
+    - nearest
+    align_corners:
+    - true
+    - true
+  - _target_: CastToTyped
+    keys: "@image_key"
+    dtype: "$torch.float32"
+  - _target_: ScaleIntensityRanged
+    keys: "@image_key"
+    a_min: -87
+    a_max: 199
+    b_min: 0
+    b_max: 1
+    clip: true
+  - _target_: CastToTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+    dtype:
+    - "$np.float16"
+    - "$np.uint8"
+  - _target_: CastToTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+    dtype:
+    - "$torch.float32"
+    - "$torch.uint8"
+  - _target_: ToTensord
+    keys:
+    - "@image_key"
+    - "@label_key"
+loss:
+  _target_: DiceCELoss
+  include_background: false
+  to_onehot_y: true
+  softmax: true
+  squared_pred: true
+  batch: true
+  smooth_nr: 0.00001
+  smooth_dr: 0.00001
+  
+dints_space:
+  _target_: monai.networks.nets.TopologySearch
+  channel_mul: 0.5
+  num_blocks: 12
+  num_depths: 4
+  use_downsample: true
+  device: "$torch.device('cuda')"
+
+network:
+  _target_: monai.networks.nets.DiNTS
+  dints_space: "@dints_space"
+  in_channels: "@input_channels"
+  num_classes: "@output_classes"
+  use_downsample: true
diff --git a/monai_pancreas_dints/configs/train.yaml b/monai_pancreas_dints/configs/train.yaml
new file mode 100644
index 00000000..fa9fa0f0
--- /dev/null
+++ b/monai_pancreas_dints/configs/train.yaml
@@ -0,0 +1,389 @@
+---
+imports:
+- "$import glob"
+- "$import json"
+- "$import os"
+- "$import ignite"
+- "$from scipy import ndimage"
+
+bundle_root: "."
+ckpt_dir: "$@bundle_root + '/models'"
+output_dir: "$@bundle_root + '/eval'"
+
+# dataset_dir: "/workspace/data/msd/Task07_Pancreas"
+# data_list_file_path: "$@bundle_root + '/configs/dataset_0.json'"
+
+# train_datalist: "$monai.data.load_decathlon_datalist(@data_list_file_path, data_list_key='training',
+#   base_dir=@dataset_dir)"
+# val_datalist: "$monai.data.load_decathlon_datalist(@data_list_file_path, data_list_key='validation',
+#   base_dir=@dataset_dir)"
+
+# device: "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
+
+arch_ckpt_path: "$@bundle_root + '/models/search_code_18590.pt'"
+arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=torch.device('cuda'))"
+
+dints_space:
+  _target_: monai.networks.nets.TopologyInstance
+  channel_mul: 1
+  num_blocks: 12
+  num_depths: 4
+  use_downsample: true
+  arch_code:
+  - "$@arch_ckpt['arch_code_a']"
+  - "$@arch_ckpt['arch_code_c']"
+  device: "$torch.device('cuda')"
+
+input_channels: 1
+output_classes: 3
+
+network_def:
+  _target_: monai.networks.nets.DiNTS
+  dints_space: "@dints_space"
+  in_channels: "@input_channels"
+  num_classes: "@output_classes"
+  use_downsample: true
+  node_a: "$@arch_ckpt['node_a']"
+network: "$@network_def.to(@device)"
+
+loss:
+  _target_: DiceCELoss
+  include_background: false
+  to_onehot_y: true
+  softmax: true
+  squared_pred: true
+  batch: true
+  smooth_nr: 1.0e-05
+  smooth_dr: 1.0e-05
+
+# optimizer:
+#   _target_: torch.optim.SGD
+#   params: "$@network.parameters()"
+#   momentum: 0.9
+#   weight_decay: 4.0e-05
+#   lr: 0.025
+
+# lr_scheduler:
+#   _target_: torch.optim.lr_scheduler.StepLR
+#   optimizer: "@optimizer"
+#   step_size: 80
+#   gamma: 0.5
+
+image_key: image
+label_key: label
+val_interval: 10
+
+# train:
+
+train_deterministic_transforms:
+- _target_: LoadImaged
+  keys:
+  - "@image_key"
+  - "@label_key"
+- _target_: EnsureChannelFirstd
+  keys:
+  - "@image_key"
+  - "@label_key"
+- _target_: Orientationd
+  keys:
+  - "@image_key"
+  - "@label_key"
+  axcodes: RAS
+- _target_: Spacingd
+  keys:
+  - "@image_key"
+  - "@label_key"
+  pixdim:
+  - 1
+  - 1
+  - 1
+  mode:
+  - bilinear
+  - nearest
+  align_corners:
+  - true
+  - true
+- _target_: CastToTyped
+  keys: "@image_key"
+  dtype: "$torch.float32"
+- _target_: ScaleIntensityRanged
+  keys: "@image_key"
+  a_min: -87
+  a_max: 199
+  b_min: 0
+  b_max: 1
+  clip: true
+- _target_: CastToTyped
+  keys:
+  - "@image_key"
+  - "@label_key"
+  dtype:
+  - "$np.float16"
+  - "$np.uint8"
+- _target_: CopyItemsd
+  keys: "@label_key"
+  times: 1
+  names:
+  - label4crop
+- _target_: Lambdad
+  keys: label4crop
+  func: "$lambda x, s=@output_classes: np.concatenate(tuple([ndimage.binary_dilation((x==_k).astype(x.dtype),
+    iterations=48).astype(float) for _k in range(s)]), axis=0)"
+  overwrite: true
+- _target_: EnsureTyped
+  keys:
+  - "@image_key"
+  - "@label_key"
+- _target_: CastToTyped
+  keys: "@image_key"
+  dtype: "$torch.float32"
+- _target_: SpatialPadd
+  keys:
+  - "@image_key"
+  - "@label_key"
+  - label4crop
+  spatial_size:
+  - 96
+  - 96
+  - 96
+  mode:
+  - reflect
+  - constant
+  - constant
+
+train_random_transforms:
+- _target_: RandCropByLabelClassesd
+  keys:
+  - "@image_key"
+  - "@label_key"
+  label_key: label4crop
+  num_classes: "@output_classes"
+  ratios: "$[1,] * @output_classes"
+  spatial_size:
+  - 96
+  - 96
+  - 96
+  num_samples: 1
+- _target_: Lambdad
+  keys: label4crop
+  func: "$lambda x: 0"
+- _target_: RandRotated
+  keys:
+  - "@image_key"
+  - "@label_key"
+  range_x: 0.3
+  range_y: 0.3
+  range_z: 0.3
+  mode:
+  - bilinear
+  - nearest
+  prob: 0.2
+- _target_: RandZoomd
+  keys:
+  - "@image_key"
+  - "@label_key"
+  min_zoom: 0.8
+  max_zoom: 1.2
+  mode:
+  - trilinear
+  - nearest
+  align_corners:
+  - true
+  -
+  prob: 0.16
+- _target_: RandGaussianSmoothd
+  keys: "@image_key"
+  sigma_x:
+  - 0.5
+  - 1.15
+  sigma_y:
+  - 0.5
+  - 1.15
+  sigma_z:
+  - 0.5
+  - 1.15
+  prob: 0.15
+- _target_: RandScaleIntensityd
+  keys: "@image_key"
+  factors: 0.3
+  prob: 0.5
+- _target_: RandShiftIntensityd
+  keys: "@image_key"
+  offsets: 0.1
+  prob: 0.5
+- _target_: RandGaussianNoised
+  keys: "@image_key"
+  std: 0.01
+  prob: 0.15
+- _target_: RandFlipd
+  keys:
+  - "@image_key"
+  - "@label_key"
+  spatial_axis: 0
+  prob: 0.5
+- _target_: RandFlipd
+  keys:
+  - "@image_key"
+  - "@label_key"
+  spatial_axis: 1
+  prob: 0.5
+- _target_: RandFlipd
+  keys:
+  - "@image_key"
+  - "@label_key"
+  spatial_axis: 2
+  prob: 0.5
+- _target_: CastToTyped
+  keys:
+  - "@image_key"
+  - "@label_key"
+  dtype:
+  - "$torch.float32"
+  - "$torch.uint8"
+- _target_: ToTensord
+  keys:
+  - "@image_key"
+  - "@label_key"
+
+train_preprocessing:
+  _target_: Compose
+  transforms: "$@train_deterministic_transforms + @train_random_transforms"
+
+# train_dataset:
+#   _target_: CacheDataset
+#   data: "@train_datalist"
+#   transform: "@train#preprocessing"
+#   cache_rate: 0.125
+#   num_workers: 4
+
+# dataloader:
+#   _target_: DataLoader
+#   dataset: "@train#dataset"
+#   batch_size: 2
+#   shuffle: true
+#   num_workers: 4
+
+inferer:
+  _target_: SimpleInferer
+
+postprocessing:
+  _target_: Compose
+  transforms:
+  - _target_: Activationsd
+    keys: pred
+    softmax: true
+  - _target_: AsDiscreted
+    keys:
+    - pred
+    - label
+    argmax:
+    - true
+    - false
+    to_onehot: "@output_classes"
+
+  handlers:
+  - _target_: LrScheduleHandler
+    lr_scheduler: "@lr_scheduler"
+    print_lr: true
+  - _target_: ValidationHandler
+    validator: "@validate#evaluator"
+    epoch_level: true
+    interval: "@val_interval"
+  - _target_: StatsHandler
+    tag_name: train_loss
+    output_transform: "$monai.handlers.from_engine(['loss'], first=True)"
+  - _target_: TensorBoardStatsHandler
+    log_dir: "@output_dir"
+    tag_name: train_loss
+    output_transform: "$monai.handlers.from_engine(['loss'], first=True)"
+
+  key_metric:
+    train_accuracy:
+      _target_: ignite.metrics.Accuracy
+      output_transform: "$monai.handlers.from_engine(['pred', 'label'])"
+
+  trainer:
+    _target_: SupervisedTrainer
+    max_epochs: 400
+    device: "@device"
+    train_data_loader: "@train#dataloader"
+    network: "@network"
+    loss_function: "@loss"
+    optimizer: "@optimizer"
+    inferer: "@train#inferer"
+    postprocessing: "@train#postprocessing"
+    key_train_metric: "@train#key_metric"
+    train_handlers: "@train#handlers"
+    amp: true
+
+# validate:
+
+val_preprocessing:
+  _target_: Compose
+  transforms: "$@train_deterministic_transforms"
+
+# val_dataset:
+#   _target_: CacheDataset
+#   data: "@val_datalist"
+#   transform: "@validate#preprocessing"
+#   cache_rate: 0.125
+
+# dataloader:
+#   _target_: DataLoader
+#   dataset: "@validate#dataset"
+#   batch_size: 1
+#   shuffle: false
+#   num_workers: 4
+
+inferer:
+  _target_: SlidingWindowInferer
+  roi_size:
+  - 96
+  - 96
+  - 96
+  sw_batch_size: 6
+  overlap: 0.625
+
+# postprocessing: "%train#postprocessing"
+
+handlers:
+- _target_: StatsHandler
+  iteration_log: false
+- _target_: TensorBoardStatsHandler
+  log_dir: "@output_dir"
+  iteration_log: false
+- _target_: CheckpointSaver
+  save_dir: "@ckpt_dir"
+  save_dict:
+    model: "@network"
+  save_key_metric: true
+  key_metric_filename: model.pt
+
+key_metric:
+  val_mean_dice:
+    _target_: MeanDice
+    include_background: false
+    output_transform: "$monai.handlers.from_engine(['pred', 'label'])"
+
+additional_metrics:
+  val_accuracy:
+    _target_: ignite.metrics.Accuracy
+    output_transform: "$monai.handlers.from_engine(['pred', 'label'])"
+
+evaluator:
+  _target_: SupervisedEvaluator
+  device: "@device"
+  val_data_loader: "@validate#dataloader"
+  network: "@network"
+  inferer: "@validate#inferer"
+  postprocessing: "@validate#postprocessing"
+  key_val_metric: "@validate#key_metric"
+  additional_metrics: "@validate#additional_metrics"
+  val_handlers: "@validate#handlers"
+  amp: true
+
+initialize:
+- "$monai.utils.set_determinism(seed=123)"
+
+run:
+- "$@train#trainer.run()"
diff --git a/monai_pancreas_dints/docs/README.md b/monai_pancreas_dints/docs/README.md
new file mode 100644
index 00000000..8ef4d171
--- /dev/null
+++ b/monai_pancreas_dints/docs/README.md
@@ -0,0 +1,194 @@
+# Model Overview
+A neural architecture search algorithm for volumetric (3D) segmentation of the pancreas and pancreatic tumor from CT image. This model is trained using the neural network model from the neural architecture search algorithm, DiNTS [1].
+
+![image](https://developer.download.nvidia.com/assets/Clara/Images/clara_pt_net_arch_search_segmentation_workflow_4-1.png)
+
+## Data
+The training dataset is the Pancreas Task from the Medical Segmentation Decathalon. Users can find more details on the datasets at http://medicaldecathlon.com/.
+
+- Target: Pancreas and pancreatic tumor
+- Modality: Portal venous phase CT
+- Size: 420 3D volumes (282 Training +139 Testing)
+- Source: Memorial Sloan Kettering Cancer Center
+- Challenge: Label unbalance with large (background), medium (pancreas) and small (tumour) structures.
+
+### Preprocessing
+The data list/split can be created with the script `scripts/prepare_datalist.py`.
+
+```
+python scripts/prepare_datalist.py --path /path-to-Task07_Pancreas/ --output configs/dataset_0.json
+```
+
+## Training configuration
+The training was performed with at least 16GB-memory GPUs.
+
+Actual Model Input: 96 x 96 x 96
+
+### Neural Architecture Search Configuration
+The neural architecture search was performed with the following:
+
+- AMP: True
+- Optimizer: SGD
+- Initial Learning Rate: 0.025
+- Loss: DiceCELoss
+
+### Optimial Architecture Training Configuration
+The training was performed with the following:
+
+- AMP: True
+- Optimizer: SGD
+- (Initial) Learning Rate: 0.025
+- Loss: DiceCELoss
+
+The segmentation of pancreas region is formulated as the voxel-wise 3-class classification. Each voxel is predicted as either foreground (pancreas body, tumour) or background. And the model is optimized with gradient descent method minimizing soft dice loss and cross-entropy loss between the predicted mask and ground truth segmentation.
+
+### Input
+One channel
+- CT image
+
+### Output
+Three channels
+- Label 2: pancreatic tumor
+- Label 1: pancreas
+- Label 0: everything else
+
+### Memory Consumption
+
+- Dataset Manager: CacheDataset
+- Data Size: 420 3D Volumes
+- Cache Rate: 1.0
+- Multi GPU (8 GPUs) - System RAM Usage: 400G
+
+### Memory Consumption Warning
+
+If you face memory issues with CacheDataset, you can either switch to a regular Dataset class or lower the caching rate `cache_rate` in the configurations within range [0, 1] to minimize the System RAM requirements.
+
+## Performance
+Dice score is used for evaluating the performance of the model. This model achieves a mean dice score of 0.62.
+
+Please note that this bundle is non-deterministic because of the trilinear interpolation used in the network. Therefore, reproducing the training process may not get exactly the same performance.
+Please refer to https://pytorch.org/docs/stable/notes/randomness.html#reproducibility for more details about reproducibility.
+
+#### Training Loss
+The loss over 3200 epochs (the bright curve is smoothed, and the dark one is the actual curve)
+
+![Training loss over 3200 epochs (the bright curve is smoothed, and the dark one is the actual curve)](https://developer.download.nvidia.com/assets/Clara/Images/clara_pt_net_arch_search_segmentation_train_4-3.png)
+
+#### Validation Dice
+The mean dice score over 3200 epochs (the bright curve is smoothed, and the dark one is the actual curve)
+
+![Validation mean dice score over 3200 epochs (the bright curve is smoothed, and the dark one is the actual curve)](https://developer.download.nvidia.com/assets/Clara/Images/clara_pt_net_arch_search_segmentation_validation_4-3.png)
+
+#### TensorRT speedup
+This bundle supports acceleration with TensorRT. The table below displays the speedup ratios observed on an A100 80G GPU.
+
+| method | torch_fp32(ms) | torch_amp(ms) | trt_fp32(ms) | trt_fp16(ms) | speedup amp | speedup fp32 | speedup fp16 | amp vs fp16|
+| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
+| model computation | 54611.72 | 19240.66 | 16104.8 | 11443.57 | 2.84 | 3.39 | 4.77 | 1.68 |
+| end2end | 133.93 | 43.41 | 35.65 | 26.63 | 3.09 | 3.76 | 5.03 | 1.63 |
+
+Where:
+- `model computation` means the speedup ratio of model's inference with a random input without preprocessing and postprocessing
+- `end2end` means run the bundle end-to-end with the TensorRT based model.
+- `torch_fp32` and `torch_amp` are for the PyTorch models with or without `amp` mode.
+- `trt_fp32` and `trt_fp16` are for the TensorRT based models converted in corresponding precision.
+- `speedup amp`, `speedup fp32` and `speedup fp16` are the speedup ratios of corresponding models versus the PyTorch float32 model
+- `amp vs fp16` is the speedup ratio between the PyTorch amp model and the TensorRT float16 based model.
+
+This result is benchmarked under:
+ - TensorRT: 8.6.1+cuda12.0
+ - Torch-TensorRT Version: 1.4.0
+ - CPU Architecture: x86-64
+ - OS: ubuntu 20.04
+ - Python version:3.8.10
+ - CUDA version: 12.1
+ - GPU models and configuration: A100 80G
+
+### Searched Architecture Visualization
+Users can install Graphviz for visualization of searched architectures (needed in [decode_plot.py](https://github.com/Project-MONAI/tutorials/blob/main/automl/DiNTS/decode_plot.py)). The edges between nodes indicate global structure, and numbers next to edges represent different operations in the cell searching space. An example of searched architecture is shown as follows:
+
+![Example of Searched Architecture](https://developer.download.nvidia.com/assets/Clara/Images/clara_pt_net_arch_search_segmentation_searched_arch_example_1.png)
+
+## MONAI Bundle Commands
+In addition to the Pythonic APIs, a few command line interfaces (CLI) are provided to interact with the bundle. The CLI supports flexible use cases, such as overriding configs at runtime and predefining arguments in a file.
+
+For more details usage instructions, visit the [MONAI Bundle Configuration Page](https://docs.monai.io/en/latest/config_syntax.html).
+
+#### Execute model searching:
+
+```
+python -m scripts.search run --config_file configs/search.yaml
+```
+
+#### Execute multi-GPU model searching (recommended):
+
+```
+torchrun --nnodes=1 --nproc_per_node=8 -m scripts.search run --config_file configs/search.yaml
+```
+
+#### Execute training:
+
+```
+python -m monai.bundle run --config_file configs/train.yaml
+```
+
+Please note that if the default dataset path is not modified with the actual path in the bundle config files, you can also override it by using `--dataset_dir`:
+
+```
+python -m monai.bundle run --config_file configs/train.yaml --dataset_dir <actual dataset path>
+```
+
+#### Override the `train` config to execute multi-GPU training:
+
+```
+torchrun --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train.yaml','configs/multi_gpu_train.yaml']"
+```
+
+#### Override the `train` config to execute evaluation with the trained model:
+
+```
+python -m monai.bundle run --config_file "['configs/train.yaml','configs/evaluate.yaml']"
+```
+
+#### Execute inference:
+
+```
+python -m monai.bundle run --config_file configs/inference.yaml
+```
+
+#### Export checkpoint for TorchScript:
+
+```
+python -m monai.bundle ckpt_export network_def --filepath models/model.ts --ckpt_file models/model.pt --meta_file configs/metadata.json --config_file configs/inference.yaml
+```
+
+#### Export checkpoint to TensorRT based models with fp32 or fp16 precision:
+
+```
+python -m monai.bundle trt_export --net_id network_def --filepath models/model_trt.ts --ckpt_file models/model.pt --meta_file configs/metadata.json --config_file configs/inference.yaml  --precision <fp32/fp16> --use_trace "True" --dynamic_batchsize "[1, 4, 8]" --converter_kwargs "{'truncate_long_and_double':True, 'torch_executed_ops': ['aten::upsample_trilinear3d']}"
+```
+
+#### Execute inference with the TensorRT model:
+
+```
+python -m monai.bundle run --config_file "['configs/inference.yaml', 'configs/inference_trt.yaml']"
+```
+
+# References
+
+[1] He, Y., Yang, D., Roth, H., Zhao, C. and Xu, D., 2021. Dints: Differentiable neural network topology search for 3d medical image segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 5841-5850).
+
+# License
+Copyright (c) MONAI Consortium
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
diff --git a/monai_pancreas_dints/docs/data_license.txt b/monai_pancreas_dints/docs/data_license.txt
new file mode 100644
index 00000000..5cffccb1
--- /dev/null
+++ b/monai_pancreas_dints/docs/data_license.txt
@@ -0,0 +1,6 @@
+Third Party Licenses
+-----------------------------------------------------------------------
+
+/*********************************************************************/
+i. Medical Segmentation Decathlon
+   http://medicaldecathlon.com/
diff --git a/monai_pancreas_dints/loops.py b/monai_pancreas_dints/loops.py
new file mode 100644
index 00000000..8921f7dc
--- /dev/null
+++ b/monai_pancreas_dints/loops.py
@@ -0,0 +1,527 @@
+
+import torch
+import numpy as np
+import torch.distributed as dist
+from torch.cuda.amp import autocast
+# from datetime import datetime
+# from scipy import ndimage
+import torch.nn.functional as F
+
+from monai.inferers import sliding_window_inference
+from monai.metrics import compute_dice
+# import yaml, time, 
+import os
+import utils
+from cycling_utils import atomic_torch_save
+from torch.utils.tensorboard import SummaryWriter
+
+def search_one_epoch(
+    model, optimizer, dints_space, arch_optimizer_a, arch_optimizer_c, 
+    train_sampler, val_sampler, scaler, train_metrics, val_metric, best_metric,
+    epoch, train_loader, loss_func, args, timer
+):
+    device = args["device"] # for convenience
+
+    decay = 0.5 ** np.sum(
+        [(epoch - args["num_epochs_warmup"]) / (args["num_epochs"] - args["num_epochs_warmup"]) > args["learning_rate_milestones"]]
+    )
+    lr = args["learning_rate"] * decay * args["world_size"]
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr
+
+    model.train()
+
+    timer.report('model.train()')
+
+    train_step = train_sampler.progress // train_loader.batch_size
+    total_steps = int(len(train_sampler) / train_loader.batch_size)
+    print(f'\nTraining / resuming epoch {epoch} from training step {train_step}\n')
+
+    for batch_data in train_loader:
+
+        inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
+        inputs_search, labels_search = inputs.detach().clone(), labels.detach().clone() # added, will this work?
+
+        timer.report('data to device')
+
+        # UPDATE MODEL
+
+        for p in model.module.weight_parameters():
+            p.requires_grad=True
+        dints_space.log_alpha_a.requires_grad = False
+        dints_space.log_alpha_c.requires_grad = False
+
+        optimizer.zero_grad()
+
+        timer.report('config model to train')
+
+        if args["amp"]:
+            with autocast():
+                outputs = model(inputs)
+                timer.report('model forward')
+                if args["output_classes"] == 2:
+                    loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+                else:
+                    loss = loss_func(outputs, labels)
+                timer.report('model loss')
+
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+            timer.report('model backward')
+        else:
+            outputs = model(inputs)
+            timer.report('model forward')
+            if args["output_classes"] == 2:
+                loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+            else:
+                loss = loss_func(outputs, labels)
+            timer.report('model loss')
+            loss.backward()
+            optimizer.step()
+            timer.report('model backward')
+
+        # Reporting and stuff
+        train_metrics.update({"model_loss": loss.item(), "inputs_seen": len(inputs)})
+
+        timer.report('model update')
+
+        # Only update space after number of warmup epochs
+        if epoch >= args["num_epochs_warmup"]:
+
+            # UPDATE SPACE
+
+            for p in model.module.weight_parameters():
+                p.requires_grad=False
+            dints_space.log_alpha_a.requires_grad = True
+            dints_space.log_alpha_c.requires_grad = True
+
+            # linear increase topology and RAM loss
+            entropy_alpha_c = torch.tensor(0.0,).to(device)
+            entropy_alpha_a = torch.tensor(0.0).to(device)
+            ram_cost_full = torch.tensor(0.0).to(device)
+            ram_cost_usage = torch.tensor(0.0).to(device)
+            ram_cost_loss = torch.tensor(0.0).to(device)
+            topology_loss = torch.tensor(0.0).to(device)
+
+            probs_a, arch_code_prob_a = dints_space.get_prob_a(child=True)
+            entropy_alpha_a = -((probs_a) * torch.log(probs_a + 1e-5)).mean()
+            sm = F.softmax(dints_space.log_alpha_c, dim=-1)
+            lsm = F.log_softmax(dints_space.log_alpha_c, dim=-1)
+            entropy_alpha_c = -(sm * lsm).mean()
+            topology_loss = dints_space.get_topology_entropy(probs_a)
+
+            ram_cost_full = dints_space.get_ram_cost_usage(inputs.shape, full=True)
+            ram_cost_usage = dints_space.get_ram_cost_usage(inputs.shape)
+            ram_cost_loss = torch.abs(args["ram_cost_factor"] - ram_cost_usage / ram_cost_full)
+
+            arch_optimizer_a.zero_grad()
+            arch_optimizer_c.zero_grad()
+
+            combination_weights = (epoch - args["num_epochs_warmup"]) / (args["num_epochs"] - args["num_epochs_warmup"])
+
+            timer.report('space combination_weights')
+
+            if args["amp"]:
+                with autocast():
+                    outputs_search = model(inputs_search)
+                    timer.report('space forward')
+                    if args["output_classes"] == 2:
+                        loss = loss_func(torch.flip(outputs_search, dims=[1]), 1 - labels_search)
+                    else:
+                        loss = loss_func(outputs_search, labels_search)
+
+                    loss += combination_weights * (
+                        (entropy_alpha_a + entropy_alpha_c) + ram_cost_loss + 0.001 * topology_loss
+                    )
+                    timer.report('space loss')
+
+                scaler.scale(loss).backward()
+                scaler.step(arch_optimizer_a)
+                scaler.step(arch_optimizer_c)
+                scaler.update()
+                timer.report('space backward')
+            else:
+                outputs_search = model(inputs_search)
+                timer.report('space forward')
+                if args["output_classes"] == 2:
+                    loss = loss_func(torch.flip(outputs_search, dims=[1]), 1 - labels_search)
+                else:
+                    loss = loss_func(outputs_search, labels_search)
+
+                loss += 1.0 * (
+                    combination_weights * (entropy_alpha_a + entropy_alpha_c) + ram_cost_loss + 0.001 * topology_loss
+                )
+                timer.report('space loss')
+
+                loss.backward()
+                arch_optimizer_a.step()
+                arch_optimizer_c.step()
+                timer.report('space backward')
+
+            # Reporting and stuff
+            train_metrics.update({"space_loss": loss.item()})
+
+            timer.report('space update')
+
+        # Batch reporting
+        train_metrics.reduce()
+        batch_model_loss = train_metrics.local["model_loss"] / train_metrics.local["inputs_seen"]
+        if "space_loss" in train_metrics.local:
+            batch_space_loss = train_metrics.local["space_loss"] / train_metrics.local["inputs_seen"]
+        else:
+            batch_space_loss = 0.0
+        print(f"EPOCH [{epoch}], BATCH [{train_step}], MODEL LOSS [{batch_model_loss:,.3f}, SPACE LOSS: [{batch_space_loss:,.3f}]")
+        train_metrics.reset_local()
+
+        timer.report('metrics reduce')
+
+        ## Checkpointing
+        print(f"Saving checkpoint at epoch {epoch} train batch {train_step}")
+        train_sampler.advance(len(inputs))
+        train_step = train_sampler.progress // train_loader.batch_size
+
+        if train_step == total_steps:
+            train_metrics.end_epoch()
+
+        if utils.is_main_process() and train_step % 1 == 0: # Checkpointing every batch
+
+            writer = SummaryWriter(log_dir=args["tboard_path"])
+            writer.add_scalar("Train/model_loss", batch_model_loss, train_step + epoch * total_steps)
+            if batch_space_loss != "NONE":
+                writer.add_scalar("Train/space_loss", batch_space_loss, train_step + epoch * total_steps)
+            writer.flush()
+            writer.close()
+
+            checkpoint = {
+                "epoch": epoch,
+                "model": model.module.state_dict(),
+                "dints": dints_space.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "arch_optimizer_a": arch_optimizer_a.state_dict(),
+                "arch_optimizer_c": arch_optimizer_c.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "val_sampler": val_sampler.state_dict(),
+                "scaler": scaler.state_dict(),
+                "train_metrics": train_metrics,
+                "val_metric": val_metric,
+                "best_metric": best_metric,
+            }
+            timer = atomic_torch_save(checkpoint, args["resume"], timer)
+
+    timer.report(f'EPOCH {epoch}')
+
+    return model, dints_space, timer, train_metrics
+
+
+def eval_search(
+        model, optimizer, dints_space, arch_optimizer_a, arch_optimizer_c, 
+        train_sampler, val_sampler, scaler, train_metrics, val_metric, best_metric,
+        epoch, val_loader, post_pred, post_label, args, timer,
+):
+    device = args["device"] # for convenience
+
+    torch.cuda.empty_cache()
+    model.eval()
+
+    timer.report('model ready to eval')
+
+    with torch.no_grad():
+
+        val_step = val_sampler.progress // val_loader.batch_size
+        total_steps = int(len(val_sampler) / val_loader.batch_size)
+        print(f'\nEvaluating / resuming epoch {epoch} from eval step {val_step}\n')
+
+        for val_data in val_loader:
+
+            val_images, val_labels = val_data["image"].to(device), val_data["label"].to(device)
+            roi_size = args["patch_size_valid"]
+            sw_batch_size = args["num_sw_batch_size"]
+
+            if args["amp"]:
+                with torch.cuda.amp.autocast():
+                    pred = sliding_window_inference(
+                        val_images, roi_size, sw_batch_size,
+                        lambda x: model(x), mode="gaussian",
+                        overlap=args["overlap_ratio"],
+                    )
+            else:
+                pred = sliding_window_inference(
+                    val_images, roi_size, sw_batch_size,
+                    lambda x: model(x), mode="gaussian",
+                    overlap=args["overlap_ratio"],
+                )
+
+            val_outputs = post_pred(pred[0, ...])
+            val_outputs = val_outputs[None, ...]
+            val_labels = post_label(val_labels[0, ...])
+            val_labels = val_labels[None, ...]
+
+            value = compute_dice(y_pred=val_outputs, y=val_labels, include_background=False)
+
+            for _c in range(args["output_classes"] - 1):
+
+                val0 = torch.nan_to_num(value[0, _c], nan=0.0)
+                val1 = 1.0 - torch.isnan(value[0, 0]).float()
+
+                val0, val1, val_metric =  val0.to(device), val1.to(device), val_metric.to(device)
+
+                print(f"val_metric.device = {val_metric.device}")
+                print(f"val0.device = {val0.device}")
+                print(f"val1.device = {val1.device}")
+
+                val_metric[2 * _c] += val0 * val1
+                val_metric[2 * _c + 1] += val1
+
+            ## Checkpointing
+            print(f"Saving checkpoint at epoch {epoch} eval batch {val_step}")
+            val_sampler.advance(len(val_images))
+            val_step = val_sampler.progress // val_loader.batch_size
+
+            if utils.is_main_process() and val_step % 1 == 0: # Checkpointing every batch
+
+                checkpoint = {
+                    "epoch": epoch,
+                    "model": model.module.state_dict(),
+                    "dints": dints_space.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "arch_optimizer_a": arch_optimizer_a.state_dict(),
+                    "arch_optimizer_c": arch_optimizer_c.state_dict(),
+                    "train_sampler": train_sampler.state_dict(),
+                    "val_sampler": val_sampler.state_dict(),
+                    "scaler": scaler.state_dict(),
+                    "train_metrics": train_metrics,
+                    "val_metric": val_metric,
+                    "best_metric": best_metric,
+                }
+                timer = atomic_torch_save(checkpoint, args["resume"], timer)
+
+            timer.report(f'eval step {val_step}')
+
+        # synchronizes all processes and reduce results
+        if torch.cuda.device_count() > 1:
+            dist.barrier()
+            val_metric = val_metric.to_dense()
+            val_metric = val_metric.to(device)
+            print(f"val_metric.device = {val_metric.device}")
+            print(f"val_metric.layout = {val_metric.layout}")
+            dist.all_reduce(val_metric, op=torch.distributed.ReduceOp.SUM)
+
+        val_metric = val_metric.tolist()
+        if utils.is_main_process():
+
+            for _c in range(args["output_classes"] - 1):
+                print("evaluation metric - class {0:d}:".format(_c + 1), val_metric[2 * _c] / val_metric[2 * _c + 1])
+            avg_metric = 0
+            for _c in range(args["output_classes"] - 1):
+                avg_metric += val_metric[2 * _c] / val_metric[2 * _c + 1]
+            avg_metric = avg_metric / float(args["output_classes"] - 1)
+            print("avg_metric", avg_metric)
+
+            if avg_metric > best_metric:
+                best_metric = avg_metric
+                # best_metric_epoch = epoch + 1
+                # best_metric_iterations = idx_iter
+
+                (node_a_d, arch_code_a_d, arch_code_c_d, arch_code_a_max_d) = dints_space.decode()
+                torch.save(
+                    {
+                        "node_a": node_a_d,
+                        "arch_code_a": arch_code_a_d,
+                        "arch_code_a_max": arch_code_a_max_d,
+                        "arch_code_c": arch_code_c_d,
+                        # "iter_num": idx_iter,
+                        "epochs": epoch + 1,
+                        "best_dsc": best_metric,
+                        # "best_path": best_metric_iterations,
+                    },
+                    os.path.join(args["arch_ckpt_path"], "search_code.pt"),
+                )
+
+                checkpoint = {
+                    "epoch": epoch,
+                    "model": model.module.state_dict(),
+                    "dints": dints_space.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "arch_optimizer_a": arch_optimizer_a.state_dict(),
+                    "arch_optimizer_c": arch_optimizer_c.state_dict(),
+                    "train_sampler": train_sampler.state_dict(),
+                    "val_sampler": val_sampler.state_dict(),
+                    "scaler": scaler.state_dict(),
+                    "train_metrics": train_metrics,
+                    "val_metric": val_metric,
+                    "best_metric": best_metric,
+                }
+                timer = atomic_torch_save(checkpoint, args["resume"], timer)
+    
+    timer.report(f'EVAL EPOCH {epoch}')
+
+    return timer
+
+
+def train_one_epoch(
+    model, optimizer, lr_scheduler,
+    train_sampler, val_sampler, scaler, train_metrics, val_metric,
+    epoch, train_loader, loss_func, args
+):
+    device = args["device"] # for convenience
+
+    model.train()
+
+    train_step = train_sampler.progress // train_loader.batch_size
+    total_steps = int(len(train_sampler) / train_loader.batch_size)
+    print(f'\nTraining / resuming epoch {epoch} from training step {train_step}\n')
+
+    for batch_data in train_loader:
+
+        inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
+
+        optimizer.zero_grad()
+
+        if args["amp"]:
+            with autocast():
+                outputs = model(inputs)
+                if args["output_classes"] == 2:
+                    loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+                else:
+                    loss = loss_func(outputs, labels)
+
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            outputs = model(inputs)
+            if args["output_classes"] == 2:
+                loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+            else:
+                loss = loss_func(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+        # Reporting and stuff
+        train_metrics.update({"model_loss": loss.item(), "inputs_seen": len(inputs)})
+        train_metrics.reduce()
+        batch_model_loss = train_metrics.local["model_loss"] / train_metrics.local["inputs_seen"]
+        print(f"EPOCH [{epoch}], BATCH [{train_step}], MODEL LOSS [{batch_model_loss:,.3f}]")
+        train_metrics.reset_local()
+
+        ## Checkpointing
+        print(f"Saving checkpoint at epoch {epoch} train batch {train_step}")
+        train_sampler.advance(len(inputs))
+        train_step = train_sampler.progress // train_loader.batch_size
+
+        if train_step == total_steps:
+            train_metrics.end_epoch()
+            lr_scheduler.step()
+
+        if utils.is_main_process() and train_step % 1 == 0: # Checkpointing every batch
+
+            writer = SummaryWriter(log_dir=args["tboard_path"])
+            writer.add_scalar("Train/model_loss", batch_model_loss, train_step + epoch * total_steps)
+            writer.flush()
+            writer.close()
+
+            checkpoint = {
+                "epoch": epoch,
+                "model": model.module.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "val_sampler": val_sampler.state_dict(),
+                "scaler": scaler.state_dict(),
+                "train_metrics": train_metrics,
+                "val_metric": val_metric,
+            }
+            timer = atomic_torch_save(checkpoint, args["resume"], timer)
+
+    return model, timer, train_metrics
+
+
+def evaluate(
+        model, optimizer, lr_scheduler,
+        train_sampler, val_sampler, scaler, train_metrics, val_metric,
+        epoch, val_loader, post_pred, post_label, args,
+):
+    device = args["device"] # for convenience
+
+    torch.cuda.empty_cache()
+    model.eval()
+
+    with torch.no_grad():
+
+        val_step = val_sampler.progress // val_loader.batch_size
+        print(f'\nEvaluating / resuming epoch {epoch} from eval step {val_step}\n')
+
+        for val_data in val_loader:
+
+            val_images, val_labels = val_data["image"].to(device), val_data["label"].to(device)
+            roi_size = args["patch_size_valid"]
+            sw_batch_size = args["num_sw_batch_size"]
+
+            if args["amp"]:
+                with torch.cuda.amp.autocast():
+                    pred = sliding_window_inference(
+                        val_images, roi_size, sw_batch_size,
+                        lambda x: model(x), mode="gaussian",
+                        overlap=args["overlap_ratio"],
+                    )
+            else:
+                pred = sliding_window_inference(
+                    val_images, roi_size, sw_batch_size,
+                    lambda x: model(x), mode="gaussian",
+                    overlap=args["overlap_ratio"],
+                )
+
+            val_outputs = post_pred(pred[0, ...])
+            val_outputs = val_outputs[None, ...]
+            val_labels = post_label(val_labels[0, ...])
+            val_labels = val_labels[None, ...]
+
+            value = compute_dice(y_pred=val_outputs, y=val_labels, include_background=False)
+
+            for _c in range(args["output_classes"] - 1):
+                val0 = torch.nan_to_num(value[0, _c], nan=0.0)
+                val1 = 1.0 - torch.isnan(value[0, 0]).float()
+                val_metric[2 * _c] += val0 * val1
+                val_metric[2 * _c + 1] += val1
+
+            ## Checkpointing
+            print(f"Saving checkpoint at epoch {epoch} eval batch {val_step}")
+            val_sampler.advance(len(val_images))
+            val_step = val_sampler.progress // val_loader.batch_size
+
+            if utils.is_main_process() and val_step % 1 == 0: # Checkpointing every batch
+
+                checkpoint = {
+                    "epoch": epoch,
+                    "model": model.module.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "lr_scheduler": lr_scheduler.state_dict(),
+                    "train_sampler": train_sampler.state_dict(),
+                    "val_sampler": val_sampler.state_dict(),
+                    "scaler": scaler.state_dict(),
+                    "train_metrics": train_metrics,
+                    "val_metric": val_metric,
+                }
+                timer = atomic_torch_save(checkpoint, args["resume"], timer)
+
+        # synchronizes all processes and reduce results
+        if torch.cuda.device_count() > 1:
+            dist.barrier()
+            dist.all_reduce(val_metric, op=torch.distributed.ReduceOp.SUM)
+
+        val_metric = val_metric.tolist()
+        if utils.is_main_process():
+
+            for _c in range(args["output_classes"] - 1):
+                print("evaluation metric - class {0:d}:".format(_c + 1), val_metric[2 * _c] / val_metric[2 * _c + 1])
+            avg_metric = 0
+            for _c in range(args["output_classes"] - 1):
+                avg_metric += val_metric[2 * _c] / val_metric[2 * _c + 1]
+            avg_metric = avg_metric / float(args["output_classes"] - 1)
+            print("avg_metric", avg_metric)
+
+            writer = SummaryWriter(log_dir=args["tboard_path"])
+            writer.add_scalar("Val/avg_metric", avg_metric, epoch)
+            writer.flush()
+            writer.close()
\ No newline at end of file
diff --git a/monai_pancreas_dints/monai_pancreas_dints.isc b/monai_pancreas_dints/monai_pancreas_dints.isc
new file mode 100644
index 00000000..e1694af0
--- /dev/null
+++ b/monai_pancreas_dints/monai_pancreas_dints.isc
@@ -0,0 +1,6 @@
+experiment_name="monai_pancreas_dints_lr"
+gpu_type="24GB VRAM GPU"
+nnodes = 11
+venv_path = "~/.venv/bin/activate"
+output_path = "~/outputs/monai_pancreas_dints"
+command="search.py --config-file configs/search.yaml --resume $OUTPUT_PATH/checkpoint.isc --tboard-path $OUTPUT_PATH/tb"
\ No newline at end of file
diff --git a/monai_pancreas_dints/prepare_datalist.py b/monai_pancreas_dints/prepare_datalist.py
new file mode 100644
index 00000000..c35657fb
--- /dev/null
+++ b/monai_pancreas_dints/prepare_datalist.py
@@ -0,0 +1,59 @@
+import argparse
+import glob
+import json
+import os
+
+import monai
+from sklearn.model_selection import train_test_split
+
+
+def produce_sample_dict(line: str):
+    return {"label": line, "image": line.replace("labelsTr", "imagesTr")}
+
+
+def produce_datalist(dataset_dir: str, train_size: int = 196):
+    """
+    This function is used to split the dataset.
+    It will produce "train_size" number of samples for training.
+    """
+
+    samples = sorted(glob.glob(os.path.join(dataset_dir, "labelsTr", "*"), recursive=True))
+    samples = [_item.replace(os.path.join(dataset_dir, "labelsTr"), "labelsTr") for _item in samples]
+    datalist = []
+    for line in samples:
+        datalist.append(produce_sample_dict(line))
+    train_list, other_list = train_test_split(datalist, train_size=train_size)
+    val_list, test_list = train_test_split(other_list, train_size=0.66)
+
+    return {"training": train_list, "validation": val_list, "testing": test_list}
+
+
+def main(args):
+    """
+    split the dataset and output the data list into a json file.
+    """
+    data_file_base_dir = args.path
+    output_json = args.output
+    # produce deterministic data splits
+    monai.utils.set_determinism(seed=123)
+    datalist = produce_datalist(dataset_dir=data_file_base_dir, train_size=args.train_size)
+    with open(output_json, "w") as f:
+        json.dump(datalist, f, ensure_ascii=True, indent=4)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument(
+        "--path",
+        type=str,
+        # default="/workspace/data/msd/Task07_Pancreas",
+        default="/mnt/Datasets/Open-Datasets/MONAI/Task07_Pancreas",
+        help="root path of MSD Task07_Pancreas dataset.",
+    )
+    parser.add_argument(
+        "--output", type=str, default="dataset_0.json", help="relative path of output datalist json file."
+    )
+    parser.add_argument("--train_size", type=int, default=196, help="number of training samples.")
+    args = parser.parse_args()
+
+    main(args)
diff --git a/monai_pancreas_dints/scripts/__init__.py b/monai_pancreas_dints/scripts/__init__.py
new file mode 100644
index 00000000..1e97f894
--- /dev/null
+++ b/monai_pancreas_dints/scripts/__init__.py
@@ -0,0 +1,10 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/monai_pancreas_dints/scripts/loops.py b/monai_pancreas_dints/scripts/loops.py
new file mode 100644
index 00000000..b1a14cec
--- /dev/null
+++ b/monai_pancreas_dints/scripts/loops.py
@@ -0,0 +1,482 @@
+
+import torch
+import numpy as np
+import torch.distributed as dist
+from torch.cuda.amp import autocast
+from datetime import datetime
+from scipy import ndimage
+import torch.nn.functional as F
+
+from monai.inferers import sliding_window_inference
+from monai.metrics import compute_dice
+import yaml, time, os
+import scripts.utils as utils
+from cycling_utils import atomic_torch_save
+from torch.utils.tensorboard import SummaryWriter
+
+def search_one_epoch(
+    model, optimizer, dints_space, arch_optimizer_a, arch_optimizer_c, 
+    train_sampler, val_sampler, scaler, train_metrics, val_metric,
+    epoch, train_loader, loss_func, args, timer
+):
+    device = args["device"] # for convenience
+
+    decay = 0.5 ** np.sum(
+        [(epoch - args["num_epochs_warmup"]) / (args["num_epochs"] - args["num_epochs_warmup"]) > args["learning_rate_milestones"]]
+    )
+    lr = args["learning_rate"] * decay * args["world_size"]
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr
+
+    model.train()
+
+    timer.report('model.train()')
+
+    train_step = train_sampler.progress // train_loader.batch_size
+    total_steps = int(len(train_sampler) / train_loader.batch_size)
+    print(f'\nTraining / resuming epoch {epoch} from training step {train_step}\n')
+
+    for batch_data in train_loader:
+
+        inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
+        inputs_search, labels_search = inputs.detach().clone(), labels.detach().clone() # added, will this work?
+
+        timer.report('data to device')
+
+        # UPDATE MODEL
+
+        for p in model.module.weight_parameters():
+            p.requires_grad=True
+        dints_space.log_alpha_a.requires_grad = False
+        dints_space.log_alpha_c.requires_grad = False
+
+        optimizer.zero_grad()
+
+        timer.report('config model to train')
+
+        if args["amp"]:
+            with autocast():
+                outputs = model(inputs)
+                if args["output_classes"] == 2:
+                    loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+                else:
+                    loss = loss_func(outputs, labels)
+
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            outputs = model(inputs)
+            if args["output_classes"] == 2:
+                loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+            else:
+                loss = loss_func(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+        # Reporting and stuff
+        train_metrics.update({"model_loss": loss.item(), "inputs_seen": len(inputs)})
+
+        timer.report('update')
+
+        # Only update space after number of warmup epochs
+        if epoch >= args["num_epochs_warmup"]:
+
+            # UPDATE SPACE
+
+            for p in model.module.weight_parameters():
+                p.requires_grad=False
+            dints_space.log_alpha_a.requires_grad = True
+            dints_space.log_alpha_c.requires_grad = True
+
+            # linear increase topology and RAM loss
+            entropy_alpha_c = torch.tensor(0.0,).to(device)
+            entropy_alpha_a = torch.tensor(0.0).to(device)
+            ram_cost_full = torch.tensor(0.0).to(device)
+            ram_cost_usage = torch.tensor(0.0).to(device)
+            ram_cost_loss = torch.tensor(0.0).to(device)
+            topology_loss = torch.tensor(0.0).to(device)
+
+            probs_a, arch_code_prob_a = dints_space.get_prob_a(child=True)
+            entropy_alpha_a = -((probs_a) * torch.log(probs_a + 1e-5)).mean()
+            sm = F.softmax(dints_space.log_alpha_c, dim=-1)
+            lsm = F.log_softmax(dints_space.log_alpha_c, dim=-1)
+            entropy_alpha_c = -(sm * lsm).mean()
+            topology_loss = dints_space.get_topology_entropy(probs_a)
+
+            ram_cost_full = dints_space.get_ram_cost_usage(inputs.shape, full=True)
+            ram_cost_usage = dints_space.get_ram_cost_usage(inputs.shape)
+            ram_cost_loss = torch.abs(args["ram_cost_factor"] - ram_cost_usage / ram_cost_full)
+
+            arch_optimizer_a.zero_grad()
+            arch_optimizer_c.zero_grad()
+
+            combination_weights = (epoch - args["num_epochs_warmup"]) / (args["num_epochs"] - args["num_epochs_warmup"])
+
+            if args["amp"]:
+                with autocast():
+                    outputs_search = model(inputs_search)
+                    if args["output_classes"] == 2:
+                        loss = loss_func(torch.flip(outputs_search, dims=[1]), 1 - labels_search)
+                    else:
+                        loss = loss_func(outputs_search, labels_search)
+
+                    loss += combination_weights * (
+                        (entropy_alpha_a + entropy_alpha_c) + ram_cost_loss + 0.001 * topology_loss
+                    )
+
+                scaler.scale(loss).backward()
+                scaler.step(arch_optimizer_a)
+                scaler.step(arch_optimizer_c)
+                scaler.update()
+            else:
+                outputs_search = model(inputs_search)
+                if args["output_classes"] == 2:
+                    loss = loss_func(torch.flip(outputs_search, dims=[1]), 1 - labels_search)
+                else:
+                    loss = loss_func(outputs_search, labels_search)
+
+                loss += 1.0 * (
+                    combination_weights * (entropy_alpha_a + entropy_alpha_c) + ram_cost_loss + 0.001 * topology_loss
+                )
+
+                loss.backward()
+                arch_optimizer_a.step()
+                arch_optimizer_c.step()
+
+            # Reporting and stuff
+            train_metrics.update({"space_loss": loss.item()})
+
+        timer.report('space update')
+
+        # Batch reporting
+        train_metrics.reduce()
+        batch_model_loss = train_metrics.local["model_loss"] / train_metrics.local["inputs_seen"]
+        if "space_loss" in train_metrics.local:
+            batch_space_loss = train_metrics.local["space_loss"] / train_metrics.local["inputs_seen"]
+        else:
+            batch_space_loss = "NONE"
+        print(f"EPOCH [{epoch}], BATCH [{train_step}], MODEL LOSS [{batch_model_loss:,.3f}, SPACE LOSS: [{batch_space_loss:,.3f}]")
+        train_metrics.reset_local()
+
+        timer.report('metrics reduce')
+
+        ## Checkpointing
+        print(f"Saving checkpoint at epoch {epoch} train batch {train_step}")
+        train_sampler.advance(len(inputs))
+        train_step = train_sampler.progress // train_loader.batch_size
+
+        if train_step == total_steps:
+            train_metrics.end_epoch()
+
+        if utils.is_main_process() and train_step % 1 == 0: # Checkpointing every batch
+
+            writer = SummaryWriter(log_dir=args["tboard_path"])
+            writer.add_scalar("Train/model_loss", batch_model_loss, train_step + epoch * total_steps)
+            if batch_space_loss != "NONE":
+                writer.add_scalar("Train/space_loss", batch_space_loss, train_step + epoch * total_steps)
+            writer.flush()
+            writer.close()
+
+            checkpoint = {
+                "epoch": epoch,
+                "model": model.module.state_dict(),
+                "dints": dints_space.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "arch_optimizer_a": arch_optimizer_a.state_dict(),
+                "arch_optimizer_c": arch_optimizer_c.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "val_sampler": val_sampler.state_dict(),
+                "scaler": scaler.state_dict(),
+                "train_metrics": train_metrics,
+                "val_metric": val_metric
+            }
+            timer = atomic_torch_save(checkpoint, args["resume"], timer)
+
+    timer.report(f'EPOCH {epoch}')
+
+    return model, dints_space, timer, train_metrics, timer
+
+
+def eval_search(
+        model, optimizer, dints_space, arch_optimizer_a, arch_optimizer_c, 
+        train_sampler, val_sampler, scaler, train_metrics, val_metric,
+        epoch, val_loader, post_pred, post_label, args, timer,
+):
+    device = args["device"] # for convenience
+
+    torch.cuda.empty_cache()
+    model.eval()
+
+    timer.report('model ready to eval')
+
+    with torch.no_grad():
+
+        val_step = val_sampler.progress // val_loader.batch_size
+        total_steps = int(len(val_sampler) / val_loader.batch_size)
+        print(f'\nEvaluating / resuming epoch {epoch} from eval step {val_step}\n')
+
+        for val_data in val_loader:
+
+            val_images, val_labels = val_data["image"].to(device), val_data["label"].to(device)
+            roi_size = args["patch_size_valid"]
+            sw_batch_size = args["num_sw_batch_size"]
+
+            if args["amp"]:
+                with torch.cuda.amp.autocast():
+                    pred = sliding_window_inference(
+                        val_images, roi_size, sw_batch_size,
+                        lambda x: model(x), mode="gaussian",
+                        overlap=args["overlap_ratio"],
+                    )
+            else:
+                pred = sliding_window_inference(
+                    val_images, roi_size, sw_batch_size,
+                    lambda x: model(x), mode="gaussian",
+                    overlap=args["overlap_ratio"],
+                )
+
+            val_outputs = post_pred(pred[0, ...])
+            val_outputs = val_outputs[None, ...]
+            val_labels = post_label(val_labels[0, ...])
+            val_labels = val_labels[None, ...]
+
+            value = compute_dice(y_pred=val_outputs, y=val_labels, include_background=False)
+
+            for _c in range(args["output_classes"] - 1):
+                val0 = torch.nan_to_num(value[0, _c], nan=0.0)
+                val1 = 1.0 - torch.isnan(value[0, 0]).float()
+                val_metric[2 * _c] += val0 * val1
+                val_metric[2 * _c + 1] += val1
+
+            ## Checkpointing
+            print(f"Saving checkpoint at epoch {epoch} eval batch {val_step}")
+            val_sampler.advance(len(val_images))
+            val_step = val_sampler.progress // val_loader.batch_size
+
+            if utils.is_main_process() and val_step % 1 == 0: # Checkpointing every batch
+
+                checkpoint = {
+                    "epoch": epoch,
+                    "model": model.module.state_dict(),
+                    "dints": dints_space.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "arch_optimizer_a": arch_optimizer_a.state_dict(),
+                    "arch_optimizer_c": arch_optimizer_c.state_dict(),
+                    "train_sampler": train_sampler.state_dict(),
+                    "val_sampler": val_sampler.state_dict(),
+                    "scaler": scaler.state_dict(),
+                    "train_metrics": train_metrics,
+                    "val_metric": val_metric
+                }
+                timer = atomic_torch_save(checkpoint, args["resume"], timer)
+
+            timer.report(f'eval step {val_step}')
+
+        # synchronizes all processes and reduce results
+        if torch.cuda.device_count() > 1:
+            dist.barrier()
+            dist.all_reduce(val_metric, op=torch.distributed.ReduceOp.SUM)
+
+        val_metric = val_metric.tolist()
+        if utils.is_main_process():
+
+            for _c in range(args["output_classes"] - 1):
+                print("evaluation metric - class {0:d}:".format(_c + 1), val_metric[2 * _c] / val_metric[2 * _c + 1])
+            avg_metric = 0
+            for _c in range(args["output_classes"] - 1):
+                avg_metric += val_metric[2 * _c] / val_metric[2 * _c + 1]
+            avg_metric = avg_metric / float(args["output_classes"] - 1)
+            print("avg_metric", avg_metric)
+
+            if avg_metric > best_metric:
+                best_metric = avg_metric
+                # best_metric_epoch = epoch + 1
+                # best_metric_iterations = idx_iter
+
+                (node_a_d, arch_code_a_d, arch_code_c_d, arch_code_a_max_d) = dints_space.decode()
+                torch.save(
+                    {
+                        "node_a": node_a_d,
+                        "arch_code_a": arch_code_a_d,
+                        "arch_code_a_max": arch_code_a_max_d,
+                        "arch_code_c": arch_code_c_d,
+                        # "iter_num": idx_iter,
+                        "epochs": epoch + 1,
+                        "best_dsc": best_metric,
+                        # "best_path": best_metric_iterations,
+                    },
+                    os.path.join(args["arch_ckpt_path"], "search_code.pt"),
+                )
+    
+    timer.report(f'EVAL EPOCH {epoch}')
+
+    return timer
+
+
+def train_one_epoch(
+    model, optimizer, lr_scheduler,
+    train_sampler, val_sampler, scaler, train_metrics, val_metric,
+    epoch, train_loader, loss_func, args
+):
+    device = args["device"] # for convenience
+
+    model.train()
+
+    train_step = train_sampler.progress // train_loader.batch_size
+    total_steps = int(len(train_sampler) / train_loader.batch_size)
+    print(f'\nTraining / resuming epoch {epoch} from training step {train_step}\n')
+
+    for batch_data in train_loader:
+
+        inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
+
+        optimizer.zero_grad()
+
+        if args["amp"]:
+            with autocast():
+                outputs = model(inputs)
+                if args["output_classes"] == 2:
+                    loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+                else:
+                    loss = loss_func(outputs, labels)
+
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            outputs = model(inputs)
+            if args["output_classes"] == 2:
+                loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+            else:
+                loss = loss_func(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+        # Reporting and stuff
+        train_metrics.update({"model_loss": loss.item(), "inputs_seen": len(inputs)})
+        train_metrics.reduce()
+        batch_model_loss = train_metrics.local["model_loss"] / train_metrics.local["inputs_seen"]
+        print(f"EPOCH [{epoch}], BATCH [{train_step}], MODEL LOSS [{batch_model_loss:,.3f}]")
+        train_metrics.reset_local()
+
+        ## Checkpointing
+        print(f"Saving checkpoint at epoch {epoch} train batch {train_step}")
+        train_sampler.advance(len(inputs))
+        train_step = train_sampler.progress // train_loader.batch_size
+
+        if train_step == total_steps:
+            train_metrics.end_epoch()
+            lr_scheduler.step()
+
+        if utils.is_main_process() and train_step % 1 == 0: # Checkpointing every batch
+
+            writer = SummaryWriter(log_dir=args["tboard_path"])
+            writer.add_scalar("Train/model_loss", batch_model_loss, train_step + epoch * total_steps)
+            writer.flush()
+            writer.close()
+
+            checkpoint = {
+                "epoch": epoch,
+                "model": model.module.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "val_sampler": val_sampler.state_dict(),
+                "scaler": scaler.state_dict(),
+                "train_metrics": train_metrics,
+                "val_metric": val_metric,
+            }
+            timer = atomic_torch_save(checkpoint, args["resume"], timer)
+
+    return model, timer, train_metrics
+
+
+def evaluate(
+        model, optimizer, lr_scheduler,
+        train_sampler, val_sampler, scaler, train_metrics, val_metric,
+        epoch, val_loader, post_pred, post_label, args,
+):
+    device = args["device"] # for convenience
+
+    torch.cuda.empty_cache()
+    model.eval()
+
+    with torch.no_grad():
+
+        val_step = val_sampler.progress // val_loader.batch_size
+        print(f'\nEvaluating / resuming epoch {epoch} from eval step {val_step}\n')
+
+        for val_data in val_loader:
+
+            val_images, val_labels = val_data["image"].to(device), val_data["label"].to(device)
+            roi_size = args["patch_size_valid"]
+            sw_batch_size = args["num_sw_batch_size"]
+
+            if args["amp"]:
+                with torch.cuda.amp.autocast():
+                    pred = sliding_window_inference(
+                        val_images, roi_size, sw_batch_size,
+                        lambda x: model(x), mode="gaussian",
+                        overlap=args["overlap_ratio"],
+                    )
+            else:
+                pred = sliding_window_inference(
+                    val_images, roi_size, sw_batch_size,
+                    lambda x: model(x), mode="gaussian",
+                    overlap=args["overlap_ratio"],
+                )
+
+            val_outputs = post_pred(pred[0, ...])
+            val_outputs = val_outputs[None, ...]
+            val_labels = post_label(val_labels[0, ...])
+            val_labels = val_labels[None, ...]
+
+            value = compute_dice(y_pred=val_outputs, y=val_labels, include_background=False)
+
+            for _c in range(args["output_classes"] - 1):
+                val0 = torch.nan_to_num(value[0, _c], nan=0.0)
+                val1 = 1.0 - torch.isnan(value[0, 0]).float()
+                val_metric[2 * _c] += val0 * val1
+                val_metric[2 * _c + 1] += val1
+
+            ## Checkpointing
+            print(f"Saving checkpoint at epoch {epoch} eval batch {val_step}")
+            val_sampler.advance(len(val_images))
+            val_step = val_sampler.progress // val_loader.batch_size
+
+            if utils.is_main_process() and val_step % 1 == 0: # Checkpointing every batch
+
+                checkpoint = {
+                    "epoch": epoch,
+                    "model": model.module.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "lr_scheduler": lr_scheduler.state_dict(),
+                    "train_sampler": train_sampler.state_dict(),
+                    "val_sampler": val_sampler.state_dict(),
+                    "scaler": scaler.state_dict(),
+                    "train_metrics": train_metrics,
+                    "val_metric": val_metric,
+                }
+                timer = atomic_torch_save(checkpoint, args["resume"], timer)
+
+        # synchronizes all processes and reduce results
+        if torch.cuda.device_count() > 1:
+            dist.barrier()
+            dist.all_reduce(val_metric, op=torch.distributed.ReduceOp.SUM)
+
+        val_metric = val_metric.tolist()
+        if utils.is_main_process():
+
+            for _c in range(args["output_classes"] - 1):
+                print("evaluation metric - class {0:d}:".format(_c + 1), val_metric[2 * _c] / val_metric[2 * _c + 1])
+            avg_metric = 0
+            for _c in range(args["output_classes"] - 1):
+                avg_metric += val_metric[2 * _c] / val_metric[2 * _c + 1]
+            avg_metric = avg_metric / float(args["output_classes"] - 1)
+            print("avg_metric", avg_metric)
+
+            writer = SummaryWriter(log_dir=args["tboard_path"])
+            writer.add_scalar("Val/avg_metric", avg_metric, epoch)
+            writer.flush()
+            writer.close()
\ No newline at end of file
diff --git a/monai_pancreas_dints/scripts/prepare_datalist.py b/monai_pancreas_dints/scripts/prepare_datalist.py
new file mode 100644
index 00000000..c35657fb
--- /dev/null
+++ b/monai_pancreas_dints/scripts/prepare_datalist.py
@@ -0,0 +1,59 @@
+import argparse
+import glob
+import json
+import os
+
+import monai
+from sklearn.model_selection import train_test_split
+
+
+def produce_sample_dict(line: str):
+    return {"label": line, "image": line.replace("labelsTr", "imagesTr")}
+
+
+def produce_datalist(dataset_dir: str, train_size: int = 196):
+    """
+    This function is used to split the dataset.
+    It will produce "train_size" number of samples for training.
+    """
+
+    samples = sorted(glob.glob(os.path.join(dataset_dir, "labelsTr", "*"), recursive=True))
+    samples = [_item.replace(os.path.join(dataset_dir, "labelsTr"), "labelsTr") for _item in samples]
+    datalist = []
+    for line in samples:
+        datalist.append(produce_sample_dict(line))
+    train_list, other_list = train_test_split(datalist, train_size=train_size)
+    val_list, test_list = train_test_split(other_list, train_size=0.66)
+
+    return {"training": train_list, "validation": val_list, "testing": test_list}
+
+
+def main(args):
+    """
+    split the dataset and output the data list into a json file.
+    """
+    data_file_base_dir = args.path
+    output_json = args.output
+    # produce deterministic data splits
+    monai.utils.set_determinism(seed=123)
+    datalist = produce_datalist(dataset_dir=data_file_base_dir, train_size=args.train_size)
+    with open(output_json, "w") as f:
+        json.dump(datalist, f, ensure_ascii=True, indent=4)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument(
+        "--path",
+        type=str,
+        # default="/workspace/data/msd/Task07_Pancreas",
+        default="/mnt/Datasets/Open-Datasets/MONAI/Task07_Pancreas",
+        help="root path of MSD Task07_Pancreas dataset.",
+    )
+    parser.add_argument(
+        "--output", type=str, default="dataset_0.json", help="relative path of output datalist json file."
+    )
+    parser.add_argument("--train_size", type=int, default=196, help="number of training samples.")
+    args = parser.parse_args()
+
+    main(args)
diff --git a/monai_pancreas_dints/scripts/search.py b/monai_pancreas_dints/scripts/search.py
new file mode 100644
index 00000000..c4f75ffc
--- /dev/null
+++ b/monai_pancreas_dints/scripts/search.py
@@ -0,0 +1,257 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from cycling_utils import TimestampedTimer
+
+timer = TimestampedTimer('importing timer')
+
+import json
+import logging
+import os
+import random
+import sys
+import time
+from datetime import datetime
+from typing import Sequence, Union
+from scipy import ndimage
+
+import monai
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+import yaml
+from monai import transforms
+from monai.bundle import ConfigParser
+from monai.networks.nets import TopologySearch, DiNTS
+from monai.losses import DiceCELoss
+from monai.data import ThreadDataLoader, partition_dataset, DataLoader
+from monai.inferers import sliding_window_inference
+from monai.metrics import compute_dice
+from monai.utils import set_determinism
+from torch.nn.parallel import DistributedDataParallel
+
+
+import argparse
+from cycling_utils import InterruptableDistributedSampler, MetricsTracker
+from scripts.loops import search_one_epoch, eval_search
+from pathlib import Path
+import scripts.utils as utils
+
+# def get_args_parser(add_help=True):
+#     parser = argparse.ArgumentParser(description="DiNTS search", add_help=add_help)
+#     parser.add_argument("--resume", type=str, help="path of checkpoint", required=True) # for checkpointing
+#     parser.add_argument("--prev-resume", default=None, help="path of previous job checkpoint for strong fail resume", dest="prev_resume") # for checkpointing
+#     parser.add_argument("--tboard-path", default=None, help="path for saving tensorboard logs", dest="tboard_path") # for checkpointing
+#     return parser
+
+timer.report('importing everything else')
+
+def run(config_file: Union[str, Sequence[str]], resume, prev_resume=None, tboard_path=None):
+    # logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+    timer = TimestampedTimer('commencing run')
+
+    parser = ConfigParser()
+    parser.read_config(config_file)
+
+    args = {
+        "start_epoch": 0,
+        "resume": resume,
+        "prev_resume": prev_resume,
+        "tboard_path": tboard_path,
+        "arch_ckpt_path": parser["arch_ckpt_path"],
+        "amp": parser["amp"],
+        "data_file_base_dir": parser["data_file_base_dir"],
+        "data_list_file_path": parser["data_list_file_path"],
+        "determ": parser["determ"],
+        "learning_rate": parser["learning_rate"],
+        "learning_rate_arch": parser["learning_rate_arch"],
+        "learning_rate_milestones": np.array(parser["learning_rate_milestones"]),
+        "num_images_per_batch": parser["num_images_per_batch"],
+        "num_epochs": parser["num_epochs"],  # around 20k iterations
+        "num_epochs_per_validation": parser["num_epochs_per_validation"],
+        "num_epochs_warmup": parser["num_epochs_warmup"],
+        "num_sw_batch_size": parser["num_sw_batch_size"],
+        "output_classes": parser["output_classes"],
+        "overlap_ratio": parser["overlap_ratio"],
+        "patch_size_valid": parser["patch_size_valid"],
+        "ram_cost_factor": parser["ram_cost_factor"],
+    }
+
+    utils.init_distributed_mode(args) # Sets args.distributed among other things
+    assert args["distributed"] # don't support cycling when not distributed for simplicity
+    device = torch.device(args["device"])
+
+    # deterministic training
+    if args["determ"]:
+        set_determinism(seed=0)
+
+    timer.report('preliminaries')
+
+    train_transforms = parser.get_parsed_content("transform_train")
+    val_transforms = parser.get_parsed_content("transform_validation")
+
+    timer.report('transforms')
+
+    with open(args["data_list_file_path"], "r") as f:
+        json_data = json.load(f)
+
+    list_train = json_data["training"]
+    list_valid = json_data["validation"]
+
+    # training data
+    files = []
+    for _i in range(len(list_train)):
+        str_img = os.path.join(args["data_file_base_dir"], list_train[_i]["image"])
+        str_seg = os.path.join(args["data_file_base_dir"], list_train[_i]["label"])
+
+        if (not os.path.exists(str_img)) or (not os.path.exists(str_seg)):
+            continue
+
+        files.append({"image": str_img, "label": str_seg})
+    train_files = files
+
+    random.shuffle(train_files)
+
+    timer.report('training files')
+
+    # validation data
+    files = []
+    for _i in range(len(list_valid)):
+        str_img = os.path.join(args["data_file_base_dir"], list_valid[_i]["image"])
+        str_seg = os.path.join(args["data_file_base_dir"], list_valid[_i]["label"])
+
+        if (not os.path.exists(str_img)) or (not os.path.exists(str_seg)):
+            continue
+
+        files.append({"image": str_img, "label": str_seg})
+    val_files = files
+
+    timer.report('validation files')
+
+    n_workers = 1
+    cache_rate = 0.0
+    train_ds = monai.data.CacheDataset(
+        data=train_files, transform=train_transforms, cache_rate=cache_rate, num_workers=n_workers
+    )
+    val_ds = monai.data.CacheDataset(data=val_files, transform=val_transforms, cache_rate=cache_rate, num_workers=n_workers)
+
+    train_sampler = InterruptableDistributedSampler(train_ds)
+    val_sampler = InterruptableDistributedSampler(val_ds)
+
+    train_loader = DataLoader(train_ds, batch_size=1, sampler=train_sampler, num_workers=1)
+    val_loader = DataLoader(val_ds, batch_size=1, sampler=val_sampler, num_workers=1)
+
+    timer.report('datasets and dataloaders')
+
+    # # TESTING
+    # timer = TimestampedTimer("testing start")
+    # for i, batch_data in enumerate(train_loader):
+    #     inputs, labels = batch_data["image"], batch_data["label"]
+    #     timer.report("batch")
+    #     inputs.size == (1, 1, 96, 96, 96), labels.size == (1, 1, 96, 96, 96)
+
+    dints_space = TopologySearch(channel_mul=0.5, num_blocks=12, num_depths=4, use_downsample=True, device=device)
+    model = DiNTS(dints_space, in_channels=1, num_classes=3, use_downsample=True)
+    loss_func = DiceCELoss(include_background=False, to_onehot_y=True, softmax=True, squared_pred=True, batch=True, smooth_nr=1e-05, smooth_dr=1e-05)
+
+    model = model.to(device)
+    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+
+    post_pred = transforms.Compose([transforms.EnsureType(), transforms.AsDiscrete(to_onehot=args["output_classes"], argmax=True)])
+    post_label = transforms.Compose([transforms.EnsureType(), transforms.AsDiscrete(to_onehot=args["output_classes"])])
+
+    timer.report('model to device')
+
+    model_without_ddp = model
+    if args["distributed"]:
+        model = DistributedDataParallel(model, device_ids=[args["gpu"]], find_unused_parameters=True)
+        model_without_ddp = model.module
+
+    # optimizers
+    optimizer = torch.optim.SGD(
+        model_without_ddp.weight_parameters(), lr=args["learning_rate"] * args["world_size"], momentum=0.9, weight_decay=0.00004
+    )
+    arch_optimizer_a = torch.optim.Adam(
+        [dints_space.log_alpha_a], lr=args["learning_rate_arch"] * args["world_size"], betas=(0.5, 0.999), weight_decay=0.0
+    )
+    arch_optimizer_c = torch.optim.Adam(
+        [dints_space.log_alpha_c], lr=args["learning_rate_arch"] * args["world_size"], betas=(0.5, 0.999), weight_decay=0.0
+    )
+
+    timer.report('model ready to train')
+
+    # amp
+    if args["amp"]:
+        from torch.cuda.amp import GradScaler
+        scaler = GradScaler()
+        if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+            print("[info] amp enabled")
+
+    # start a typical PyTorch training
+    val_interval = args["num_epochs_per_validation"]
+
+    # Init metric trackers
+    train_metrics = MetricsTracker()
+    val_metric = torch.zeros((args["output_classes"] - 1) * 2, dtype=torch.float, device=device)
+
+    timer.report('metrics setup')
+
+    # RETRIEVE CHECKPOINT
+    Path(args["resume"]).parent.mkdir(parents=True, exist_ok=True)
+
+    checkpoint = None
+    if args["resume"] and os.path.isfile(args["resume"]): # If we're resuming...
+        checkpoint = torch.load(args["resume"], map_location="cpu")
+    elif args["prev_resume"] and os.path.isfile(args["prev_resume"]):
+        checkpoint = torch.load(args["prev_resume"], map_location="cpu")
+    if checkpoint is not None:
+        args["start_epoch"] = checkpoint["epoch"]
+        model_without_ddp.load_state_dict(checkpoint["model"])
+        dints_space.load_state_dict(checkpoint["dints"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        arch_optimizer_a.load_state_dict(checkpoint["arch_optimizer_a"])
+        arch_optimizer_c.load_state_dict(checkpoint["arch_optimizer_c"])
+        train_sampler.load_state_dict(checkpoint["train_sampler"])
+        val_sampler.load_state_dict(checkpoint["val_sampler"])
+        scaler.load_state_dict(checkpoint["scaler"])
+        train_metrics = checkpoint["train_metrics"]
+        val_metric = checkpoint["val_metric"]
+        val_metric.to(device)
+
+    timer.report('obtain checkpoint')
+
+    for epoch in range(args["start_epoch"], args["num_epochs"]):
+
+        print('\n')
+        print(f"EPOCH :: {epoch} / {args['num_epochs']}")
+        print('\n')
+
+        with train_sampler.in_epoch(epoch):
+            timer = TimestampedTimer("Start training")
+
+            model, dints_space, timer, train_metrics = search_one_epoch(
+                model, optimizer, dints_space, arch_optimizer_a, arch_optimizer_c, 
+                train_sampler, val_sampler, scaler, train_metrics, val_metric,
+                epoch, train_loader, loss_func, args, timer
+            )
+            timer.report(f'searching space for epoch {epoch}')
+
+            if (epoch + 1) % val_interval == 0 or (epoch + 1) == args["num_epochs"]:
+
+                with val_sampler.in_epoch(epoch):
+                    timer = TimestampedTimer("Start evaluation")
+
+                    timer = eval_search(
+                        model, optimizer, dints_space, arch_optimizer_a, arch_optimizer_c, 
+                        train_sampler, val_sampler, scaler, train_metrics, val_metric,
+                        epoch, val_loader, post_pred, post_label, args, timer
+                    )
+                    timer.report(f'evaluating search for epoch {epoch}')
diff --git a/monai_pancreas_dints/scripts/train.py b/monai_pancreas_dints/scripts/train.py
new file mode 100644
index 00000000..a27cef96
--- /dev/null
+++ b/monai_pancreas_dints/scripts/train.py
@@ -0,0 +1,173 @@
+from cycling_utils import TimestampedTimer
+
+timer = TimestampedTimer()
+timer.report('importing Timer')
+
+import json
+import logging
+import os
+import random
+import sys
+import time
+from datetime import datetime
+from typing import Sequence, Union
+from scipy import ndimage
+
+import monai
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+import yaml
+from monai import transforms
+from monai.bundle import ConfigParser
+from monai.networks.nets import TopologyInstance, DiNTS
+from monai.losses import DiceCELoss
+from monai.data import ThreadDataLoader, partition_dataset, DataLoader, load_decathlon_datalist
+from monai.inferers import sliding_window_inference
+from monai.metrics import compute_dice
+from monai.utils import set_determinism
+
+from torch.nn.parallel import DistributedDataParallel
+
+
+import argparse
+from cycling_utils import InterruptableDistributedSampler, MetricsTracker
+from loops import train_one_epoch, evaluate
+from pathlib import Path
+import utils
+
+# def get_args_parser(add_help=True):
+#     parser = argparse.ArgumentParser(description="DiNTS train", add_help=add_help)
+#     parser.add_argument("--resume", type=str, help="path of checkpoint", required=True) # for checkpointing
+#     parser.add_argument("--prev-resume", default=None, help="path of previous job checkpoint for strong fail resume", dest="prev_resume") # for checkpointing
+#     parser.add_argument("--tboard-path", default=None, help="path for saving tensorboard logs", dest="tboard_path") # for checkpointing
+#     parser.add_argument("--data_list_file_path", default=None, help="path for retreiving pre-prepared data list", dest="data_list_file_path")
+#     return parser
+
+def run(config_file: Union[str, Sequence[str]], resume=None, prev_resume=None, tboard_path=None):
+    # logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+
+    parser = ConfigParser()
+    parser.read_config(config_file)
+
+    args = {
+        "start_epoch": 0,
+        "resume": resume,
+        "prev_resume": prev_resume,
+        "tboard_path": tboard_path,
+        "arch_ckpt_path": "/models/search_code.pt",
+        "num_epochs_per_validation": 10,
+
+        "learning_rate": 0.025,
+        "data_list_file_path": "/configs/dataset_0.json",
+        "dataset_dir": "/mnt/Datasets/Open-Datasets/MONAI/Task07_Pancreas",
+    }
+
+    utils.init_distributed_mode(args) # Sets args.distributed among other things
+    assert args["distributed"] # don't support cycling when not distributed for simplicity
+    device = torch.device(args["device"])
+
+    train_datalist = load_decathlon_datalist(args["data_list_file_path"], data_list_key='training', base_dir=args["dataset_dir"])
+    val_datalist = load_decathlon_datalist(args["data_list_file_path"], data_list_key='validation', base_dir=args["dataset_dir"])
+
+    train_preprocessing = parser.get_parsed_content("train_preprocessing")
+    val_preprocessing = parser.get_parsed_content("val_preprocessing")
+    postprocessing = parser.get_parsed_content("postprocessing")
+
+    n_workers = 1
+    cache_rate = 0.0
+    train_ds = monai.data.CacheDataset(data=train_datalist, transform=train_preprocessing, cache_rate=cache_rate, num_workers=n_workers)
+    val_ds = monai.data.CacheDataset(data=val_datalist, transform=val_preprocessing, cache_rate=cache_rate, num_workers=n_workers)
+
+    train_sampler = InterruptableDistributedSampler(train_ds)
+    val_sampler = InterruptableDistributedSampler(val_ds)
+
+    train_loader = DataLoader(train_ds, batch_size=1, sampler=train_sampler, num_workers=1)
+    val_loader = DataLoader(val_ds, batch_size=1, sampler=val_sampler, num_workers=1)
+
+    arch_ckpt = torch.load(args["arch_ckpt_path"], map_location=device)
+    dints_space = TopologyInstance(arch_code=[arch_ckpt['arch_code_a'], arch_ckpt['arch_code_c']], channel_mul=1.0, num_blocks=12, num_depths=4, use_downsample=True, device=device)
+    model = DiNTS(dints_space, in_channels=1, num_classes=3, use_downsample=True, node_a=arch_ckpt['node_a'])
+    model = model.to(device)
+    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+    loss_func = DiceCELoss(include_background=False, to_onehot_y=True, softmax=True, squared_pred=True, batch=True, smooth_nr=1e-05, smooth_dr=1e-05)
+
+    post_pred = transforms.Compose([transforms.Activationsd(softmax=True), transforms.AsDiscrete(to_onehot=args["output_classes"], argmax=True)])
+    post_label = transforms.Compose([transforms.Activationsd(softmax=False), transforms.AsDiscrete(to_onehot=args["output_classes"], argmax=False)])
+
+    model_without_ddp = model
+    if args["distributed"]:
+        model = DistributedDataParallel(model, device_ids=[args["gpu"]], find_unused_parameters=True)
+        model_without_ddp = model.module
+
+    optimizer = torch.optim.SGD(
+        model_without_ddp.weight_parameters(), lr=args["learning_rate"] * args["world_size"], momentum=0.9, weight_decay=0.00004
+    )
+    dints_space.log_alpha_a.requires_grad = False
+    dints_space.log_alpha_c.requires_grad = False
+
+    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, milestones=[80 * args["world_size"]], gamma=0.5)
+
+    # amp
+    if args["amp"]:
+        from torch.cuda.amp import GradScaler
+        scaler = GradScaler()
+        if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+            print("[info] amp enabled")
+
+    val_interval = args["num_epochs_per_validation"]
+
+    # Init metric trackers
+    train_metrics = MetricsTracker()
+    val_metric = torch.zeros((args["output_classes"] - 1) * 2, dtype=torch.float, device=device)
+
+    # RETRIEVE CHECKPOINT
+    Path(args["resume"]).parent.mkdir(parents=True, exist_ok=True)
+
+    checkpoint = None
+    if args["resume"] and os.path.isfile(args["resume"]): # If we're resuming...
+        checkpoint = torch.load(args["resume"], map_location="cpu")
+    elif args["prev_resume"] and os.path.isfile(args["prev_resume"]):
+        checkpoint = torch.load(args["prev_resume"], map_location="cpu")
+
+    if checkpoint is not None:
+        args["start_epoch"] = checkpoint["epoch"]
+        model_without_ddp.load_state_dict(checkpoint["model"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
+        train_sampler.load_state_dict(checkpoint["train_sampler"])
+        val_sampler.load_state_dict(checkpoint["val_sampler"])
+        scaler.load_state_dict(checkpoint["scaler"])
+        train_metrics = checkpoint["train_metrics"]
+        val_metric = checkpoint["val_metric"]
+        val_metric.to(device)
+
+    for epoch in range(args["start_epoch"], args["num_epochs"]):
+
+        print('\n')
+        print(f"EPOCH :: {epoch}")
+        print('\n')
+
+        with train_sampler.in_epoch(epoch):
+            timer = TimestampedTimer("Start training")
+
+            model, dints_space, timer, train_metrics, val_metric = train_one_epoch(
+                model, optimizer, lr_scheduler,
+                train_sampler, val_sampler, scaler, train_metrics, val_metric,
+                epoch, train_loader, loss_func, args
+            )
+            timer.report(f'training for epoch {epoch}')
+
+            if (epoch + 1) % val_interval == 0 or (epoch + 1) == args["num_epochs"]:
+
+                with val_sampler.in_epoch(epoch):
+                    timer = TimestampedTimer("Start evaluation")
+
+                    timer = evaluate(
+                        model, optimizer, lr_scheduler,
+                        train_sampler, val_sampler, scaler, train_metrics, val_metric,
+                        epoch, val_loader, post_pred, post_label, args,
+                    )
+                    timer.report(f'evaluating for epoch {epoch}')
+
diff --git a/monai_pancreas_dints/scripts/utils.py b/monai_pancreas_dints/scripts/utils.py
new file mode 100644
index 00000000..581df3fe
--- /dev/null
+++ b/monai_pancreas_dints/scripts/utils.py
@@ -0,0 +1,72 @@
+import torch, os, errno
+import torch.distributed as dist
+
+def mkdir(path):
+    try:
+        os.makedirs(path)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+def init_distributed_mode(args):
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        args["rank"] = int(os.environ["RANK"])
+        args["world_size"] = int(os.environ["WORLD_SIZE"])
+        args["gpu"] = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args["rank"] = int(os.environ["SLURM_PROCID"])
+        args["gpu"] = args["rank"] % torch.cuda.device_count()
+    else:
+        print("Not using distributed mode")
+        args["distributed"] = False
+        return
+
+    args["distributed"] = True
+
+    torch.cuda.set_device(args["gpu"])
+    args["dist_backend"] = "nccl"
+    print(f"| distributed init (rank {args['rank']}): {args['dist_url']}", flush=True)
+    torch.distributed.init_process_group(
+        backend=args["dist_backend"], init_method=args["dist_url"], world_size=args["world_size"], rank=args["rank"]
+    )
+    torch.distributed.barrier()
+    setup_for_distributed(args["rank"] == 0)
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
\ No newline at end of file
diff --git a/monai_pancreas_dints/search.py b/monai_pancreas_dints/search.py
new file mode 100644
index 00000000..89b87bdb
--- /dev/null
+++ b/monai_pancreas_dints/search.py
@@ -0,0 +1,266 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from cycling_utils import TimestampedTimer
+
+timer = TimestampedTimer('importing timer')
+
+import json
+# import logging
+import os
+import random
+# import sys
+# import time
+# from datetime import datetime
+# from typing import Sequence, Union
+# from scipy import ndimage
+
+import monai
+import numpy as np
+import torch
+import torch.distributed as dist
+# import torch.nn.functional as F
+# import yaml
+from monai import transforms
+from monai.bundle import ConfigParser
+from monai.networks.nets import TopologySearch, DiNTS
+from monai.losses import DiceCELoss
+# from monai.data import ThreadDataLoader, partition_dataset, 
+from monai.data import DataLoader
+# from monai.inferers import sliding_window_inference
+# from monai.metrics import compute_dice
+from monai.utils import set_determinism
+from torch.nn.parallel import DistributedDataParallel
+
+from cycling_utils import InterruptableDistributedSampler, MetricsTracker
+from loops import search_one_epoch, eval_search
+from pathlib import Path
+import utils
+
+def get_args_parser(add_help=True):
+    import argparse
+    parser = argparse.ArgumentParser(description="DiNTS search", add_help=add_help)
+    parser.add_argument("--config-file", type=str, help="config file", required=True, dest="config_file") # for checkpointing
+    parser.add_argument("--resume", type=str, help="path of checkpoint", required=True) # for checkpointing
+    parser.add_argument("--prev-resume", default=None, help="path of previous job checkpoint for strong fail resume", dest="prev_resume") # for checkpointing
+    parser.add_argument("--tboard-path", default=None, help="path for saving tensorboard logs", dest="tboard_path") # for checkpointing
+    return parser
+
+timer.report('importing everything else')
+
+def main(args, timer):
+    # logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+    timer = TimestampedTimer('commencing run')
+
+    parser = ConfigParser()
+    parser.read_config(args.config_file)
+
+    args = {
+        "start_epoch": 0,
+        "resume": args.resume,
+        "prev_resume": args.prev_resume,
+        "tboard_path": args.tboard_path,
+        "device": "cuda",
+        "dist_url": "env://",
+        "arch_ckpt_path": parser["arch_ckpt_path"],
+        "amp": parser["amp"],
+        "data_file_base_dir": parser["data_file_base_dir"],
+        "data_list_file_path": parser["data_list_file_path"],
+        "determ": parser["determ"],
+        "learning_rate": parser["learning_rate"],
+        "learning_rate_arch": parser["learning_rate_arch"],
+        "learning_rate_milestones": np.array(parser["learning_rate_milestones"]),
+        "num_images_per_batch": parser["num_images_per_batch"],
+        "num_epochs": parser["num_epochs"],  # around 20k iterations
+        "num_epochs_per_validation": parser["num_epochs_per_validation"],
+        "num_epochs_warmup": parser["num_epochs_warmup"],
+        "num_sw_batch_size": parser["num_sw_batch_size"],
+        "output_classes": parser["output_classes"],
+        "overlap_ratio": parser["overlap_ratio"],
+        "patch_size_valid": parser["patch_size_valid"],
+        "ram_cost_factor": parser["ram_cost_factor"],
+    }
+
+    utils.init_distributed_mode(args) # Sets args.distributed among other things
+    assert args["distributed"] # don't support cycling when not distributed for simplicity
+    device = torch.device(args["device"])
+
+    # deterministic training
+    if args["determ"]:
+        set_determinism(seed=0)
+
+    timer.report('preliminaries')
+
+    train_transforms = parser.get_parsed_content("transform_train")
+    val_transforms = parser.get_parsed_content("transform_validation")
+
+    timer.report('transforms')
+
+    with open(args["data_list_file_path"], "r") as f:
+        json_data = json.load(f)
+
+    list_train = json_data["training"]
+    list_valid = json_data["validation"]
+
+    # training data
+    files = []
+    for _i in range(len(list_train)):
+        str_img = os.path.join(args["data_file_base_dir"], list_train[_i]["image"])
+        str_seg = os.path.join(args["data_file_base_dir"], list_train[_i]["label"])
+
+        if (not os.path.exists(str_img)) or (not os.path.exists(str_seg)):
+            continue
+
+        files.append({"image": str_img, "label": str_seg})
+    train_files = files
+
+    random.shuffle(train_files)
+
+    timer.report('training files')
+
+    # validation data
+    files = []
+    for _i in range(len(list_valid)):
+        str_img = os.path.join(args["data_file_base_dir"], list_valid[_i]["image"])
+        str_seg = os.path.join(args["data_file_base_dir"], list_valid[_i]["label"])
+
+        if (not os.path.exists(str_img)) or (not os.path.exists(str_seg)):
+            continue
+
+        files.append({"image": str_img, "label": str_seg})
+    val_files = files
+
+    timer.report('validation files')
+
+    n_workers = 1
+    cache_rate = 0.0
+    train_ds = monai.data.CacheDataset(
+        data=train_files, transform=train_transforms, cache_rate=cache_rate, num_workers=n_workers
+    )
+    val_ds = monai.data.CacheDataset(data=val_files, transform=val_transforms, cache_rate=cache_rate, num_workers=n_workers)
+
+    train_sampler = InterruptableDistributedSampler(train_ds)
+    val_sampler = InterruptableDistributedSampler(val_ds)
+
+    train_loader = DataLoader(train_ds, batch_size=1, sampler=train_sampler, num_workers=1)
+    val_loader = DataLoader(val_ds, batch_size=1, sampler=val_sampler, num_workers=1)
+
+    timer.report('datasets and dataloaders')
+
+    # # TESTING
+    # timer = TimestampedTimer("testing start")
+    # for i, batch_data in enumerate(train_loader):
+    #     inputs, labels = batch_data["image"], batch_data["label"]
+    #     timer.report("batch")
+    #     inputs.size == (1, 1, 96, 96, 96), labels.size == (1, 1, 96, 96, 96)
+
+    dints_space = TopologySearch(channel_mul=0.5, num_blocks=12, num_depths=4, use_downsample=True, device=device)
+    model = DiNTS(dints_space, in_channels=1, num_classes=3, use_downsample=True)
+    loss_func = DiceCELoss(include_background=False, to_onehot_y=True, softmax=True, squared_pred=True, batch=True, smooth_nr=1e-05, smooth_dr=1e-05)
+
+    model = model.to(device)
+    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+
+    post_pred = transforms.Compose([transforms.EnsureType(), transforms.AsDiscrete(to_onehot=args["output_classes"], argmax=True)])
+    post_label = transforms.Compose([transforms.EnsureType(), transforms.AsDiscrete(to_onehot=args["output_classes"])])
+
+    timer.report('model to device')
+
+    model_without_ddp = model
+    if args["distributed"]:
+        model = DistributedDataParallel(model, device_ids=[args["gpu"]], find_unused_parameters=True)
+        model_without_ddp = model.module
+
+    # optimizers
+    optimizer = torch.optim.SGD(
+        model_without_ddp.weight_parameters(), lr=args["learning_rate"] * args["world_size"], momentum=0.9, weight_decay=0.00004
+    )
+    arch_optimizer_a = torch.optim.Adam(
+        [dints_space.log_alpha_a], lr=args["learning_rate_arch"] * args["world_size"], betas=(0.5, 0.999), weight_decay=0.0
+    )
+    arch_optimizer_c = torch.optim.Adam(
+        [dints_space.log_alpha_c], lr=args["learning_rate_arch"] * args["world_size"], betas=(0.5, 0.999), weight_decay=0.0
+    )
+
+    timer.report('model ready to train')
+
+    # amp
+    if args["amp"]:
+        from torch.cuda.amp import GradScaler
+        scaler = GradScaler()
+        if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+            print("[info] amp enabled")
+
+    # start a typical PyTorch training
+    val_interval = args["num_epochs_per_validation"]
+
+    # Init metric trackers
+    train_metrics = MetricsTracker()
+    val_metric = torch.zeros((args["output_classes"] - 1) * 2, dtype=torch.float, device=device)
+    best_metric = 0
+
+    timer.report('metrics setup')
+
+    # RETRIEVE CHECKPOINT
+    Path(args["resume"]).parent.mkdir(parents=True, exist_ok=True)
+
+    checkpoint = None
+    if args["resume"] and os.path.isfile(args["resume"]): # If we're resuming...
+        checkpoint = torch.load(args["resume"], map_location="cpu")
+    elif args["prev_resume"] and os.path.isfile(args["prev_resume"]):
+        checkpoint = torch.load(args["prev_resume"], map_location="cpu")
+    if checkpoint is not None:
+        args["start_epoch"] = checkpoint["epoch"]
+        model_without_ddp.load_state_dict(checkpoint["model"])
+        dints_space.load_state_dict(checkpoint["dints"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        arch_optimizer_a.load_state_dict(checkpoint["arch_optimizer_a"])
+        arch_optimizer_c.load_state_dict(checkpoint["arch_optimizer_c"])
+        train_sampler.load_state_dict(checkpoint["train_sampler"])
+        val_sampler.load_state_dict(checkpoint["val_sampler"])
+        scaler.load_state_dict(checkpoint["scaler"])
+        train_metrics = checkpoint["train_metrics"]
+        val_metric = checkpoint["val_metric"]
+        best_metric = checkpoint["best_metric"]
+        val_metric.to(device)
+
+    timer.report('obtain checkpoint')
+
+    for epoch in range(args["start_epoch"], args["num_epochs"]):
+
+        print('\n')
+        print(f"EPOCH :: {epoch} / {args['num_epochs']}")
+        print('\n')
+
+        with train_sampler.in_epoch(epoch):
+            timer = TimestampedTimer("Start training")
+
+            model, dints_space, timer, train_metrics = search_one_epoch(
+                model, optimizer, dints_space, arch_optimizer_a, arch_optimizer_c, 
+                train_sampler, val_sampler, scaler, train_metrics, val_metric, best_metric,
+                epoch, train_loader, loss_func, args, timer
+            )
+            timer.report(f'searching space for epoch {epoch}')
+
+            if (epoch + 1) % val_interval == 0 or (epoch + 1) == args["num_epochs"]:
+
+                with val_sampler.in_epoch(epoch):
+                    timer = TimestampedTimer("Start evaluation")
+
+                    timer = eval_search(
+                        model, optimizer, dints_space, arch_optimizer_a, arch_optimizer_c, 
+                        train_sampler, val_sampler, scaler, train_metrics, val_metric, best_metric,
+                        epoch, val_loader, post_pred, post_label, args, timer
+                    )
+                    timer.report(f'evaluating search for epoch {epoch}')
+
+if __name__ == "__main__":
+    args = get_args_parser().parse_args()
+    main(args, timer)
diff --git a/monai_pancreas_dints/train.py b/monai_pancreas_dints/train.py
new file mode 100644
index 00000000..a27cef96
--- /dev/null
+++ b/monai_pancreas_dints/train.py
@@ -0,0 +1,173 @@
+from cycling_utils import TimestampedTimer
+
+timer = TimestampedTimer()
+timer.report('importing Timer')
+
+import json
+import logging
+import os
+import random
+import sys
+import time
+from datetime import datetime
+from typing import Sequence, Union
+from scipy import ndimage
+
+import monai
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+import yaml
+from monai import transforms
+from monai.bundle import ConfigParser
+from monai.networks.nets import TopologyInstance, DiNTS
+from monai.losses import DiceCELoss
+from monai.data import ThreadDataLoader, partition_dataset, DataLoader, load_decathlon_datalist
+from monai.inferers import sliding_window_inference
+from monai.metrics import compute_dice
+from monai.utils import set_determinism
+
+from torch.nn.parallel import DistributedDataParallel
+
+
+import argparse
+from cycling_utils import InterruptableDistributedSampler, MetricsTracker
+from loops import train_one_epoch, evaluate
+from pathlib import Path
+import utils
+
+# def get_args_parser(add_help=True):
+#     parser = argparse.ArgumentParser(description="DiNTS train", add_help=add_help)
+#     parser.add_argument("--resume", type=str, help="path of checkpoint", required=True) # for checkpointing
+#     parser.add_argument("--prev-resume", default=None, help="path of previous job checkpoint for strong fail resume", dest="prev_resume") # for checkpointing
+#     parser.add_argument("--tboard-path", default=None, help="path for saving tensorboard logs", dest="tboard_path") # for checkpointing
+#     parser.add_argument("--data_list_file_path", default=None, help="path for retreiving pre-prepared data list", dest="data_list_file_path")
+#     return parser
+
+def run(config_file: Union[str, Sequence[str]], resume=None, prev_resume=None, tboard_path=None):
+    # logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+
+    parser = ConfigParser()
+    parser.read_config(config_file)
+
+    args = {
+        "start_epoch": 0,
+        "resume": resume,
+        "prev_resume": prev_resume,
+        "tboard_path": tboard_path,
+        "arch_ckpt_path": "/models/search_code.pt",
+        "num_epochs_per_validation": 10,
+
+        "learning_rate": 0.025,
+        "data_list_file_path": "/configs/dataset_0.json",
+        "dataset_dir": "/mnt/Datasets/Open-Datasets/MONAI/Task07_Pancreas",
+    }
+
+    utils.init_distributed_mode(args) # Sets args.distributed among other things
+    assert args["distributed"] # don't support cycling when not distributed for simplicity
+    device = torch.device(args["device"])
+
+    train_datalist = load_decathlon_datalist(args["data_list_file_path"], data_list_key='training', base_dir=args["dataset_dir"])
+    val_datalist = load_decathlon_datalist(args["data_list_file_path"], data_list_key='validation', base_dir=args["dataset_dir"])
+
+    train_preprocessing = parser.get_parsed_content("train_preprocessing")
+    val_preprocessing = parser.get_parsed_content("val_preprocessing")
+    postprocessing = parser.get_parsed_content("postprocessing")
+
+    n_workers = 1
+    cache_rate = 0.0
+    train_ds = monai.data.CacheDataset(data=train_datalist, transform=train_preprocessing, cache_rate=cache_rate, num_workers=n_workers)
+    val_ds = monai.data.CacheDataset(data=val_datalist, transform=val_preprocessing, cache_rate=cache_rate, num_workers=n_workers)
+
+    train_sampler = InterruptableDistributedSampler(train_ds)
+    val_sampler = InterruptableDistributedSampler(val_ds)
+
+    train_loader = DataLoader(train_ds, batch_size=1, sampler=train_sampler, num_workers=1)
+    val_loader = DataLoader(val_ds, batch_size=1, sampler=val_sampler, num_workers=1)
+
+    arch_ckpt = torch.load(args["arch_ckpt_path"], map_location=device)
+    dints_space = TopologyInstance(arch_code=[arch_ckpt['arch_code_a'], arch_ckpt['arch_code_c']], channel_mul=1.0, num_blocks=12, num_depths=4, use_downsample=True, device=device)
+    model = DiNTS(dints_space, in_channels=1, num_classes=3, use_downsample=True, node_a=arch_ckpt['node_a'])
+    model = model.to(device)
+    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+    loss_func = DiceCELoss(include_background=False, to_onehot_y=True, softmax=True, squared_pred=True, batch=True, smooth_nr=1e-05, smooth_dr=1e-05)
+
+    post_pred = transforms.Compose([transforms.Activationsd(softmax=True), transforms.AsDiscrete(to_onehot=args["output_classes"], argmax=True)])
+    post_label = transforms.Compose([transforms.Activationsd(softmax=False), transforms.AsDiscrete(to_onehot=args["output_classes"], argmax=False)])
+
+    model_without_ddp = model
+    if args["distributed"]:
+        model = DistributedDataParallel(model, device_ids=[args["gpu"]], find_unused_parameters=True)
+        model_without_ddp = model.module
+
+    optimizer = torch.optim.SGD(
+        model_without_ddp.weight_parameters(), lr=args["learning_rate"] * args["world_size"], momentum=0.9, weight_decay=0.00004
+    )
+    dints_space.log_alpha_a.requires_grad = False
+    dints_space.log_alpha_c.requires_grad = False
+
+    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, milestones=[80 * args["world_size"]], gamma=0.5)
+
+    # amp
+    if args["amp"]:
+        from torch.cuda.amp import GradScaler
+        scaler = GradScaler()
+        if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+            print("[info] amp enabled")
+
+    val_interval = args["num_epochs_per_validation"]
+
+    # Init metric trackers
+    train_metrics = MetricsTracker()
+    val_metric = torch.zeros((args["output_classes"] - 1) * 2, dtype=torch.float, device=device)
+
+    # RETRIEVE CHECKPOINT
+    Path(args["resume"]).parent.mkdir(parents=True, exist_ok=True)
+
+    checkpoint = None
+    if args["resume"] and os.path.isfile(args["resume"]): # If we're resuming...
+        checkpoint = torch.load(args["resume"], map_location="cpu")
+    elif args["prev_resume"] and os.path.isfile(args["prev_resume"]):
+        checkpoint = torch.load(args["prev_resume"], map_location="cpu")
+
+    if checkpoint is not None:
+        args["start_epoch"] = checkpoint["epoch"]
+        model_without_ddp.load_state_dict(checkpoint["model"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
+        train_sampler.load_state_dict(checkpoint["train_sampler"])
+        val_sampler.load_state_dict(checkpoint["val_sampler"])
+        scaler.load_state_dict(checkpoint["scaler"])
+        train_metrics = checkpoint["train_metrics"]
+        val_metric = checkpoint["val_metric"]
+        val_metric.to(device)
+
+    for epoch in range(args["start_epoch"], args["num_epochs"]):
+
+        print('\n')
+        print(f"EPOCH :: {epoch}")
+        print('\n')
+
+        with train_sampler.in_epoch(epoch):
+            timer = TimestampedTimer("Start training")
+
+            model, dints_space, timer, train_metrics, val_metric = train_one_epoch(
+                model, optimizer, lr_scheduler,
+                train_sampler, val_sampler, scaler, train_metrics, val_metric,
+                epoch, train_loader, loss_func, args
+            )
+            timer.report(f'training for epoch {epoch}')
+
+            if (epoch + 1) % val_interval == 0 or (epoch + 1) == args["num_epochs"]:
+
+                with val_sampler.in_epoch(epoch):
+                    timer = TimestampedTimer("Start evaluation")
+
+                    timer = evaluate(
+                        model, optimizer, lr_scheduler,
+                        train_sampler, val_sampler, scaler, train_metrics, val_metric,
+                        epoch, val_loader, post_pred, post_label, args,
+                    )
+                    timer.report(f'evaluating for epoch {epoch}')
+
diff --git a/monai_pancreas_dints/utils.py b/monai_pancreas_dints/utils.py
new file mode 100644
index 00000000..2149b1e2
--- /dev/null
+++ b/monai_pancreas_dints/utils.py
@@ -0,0 +1,72 @@
+import torch, os, errno
+import torch.distributed as dist
+
+def mkdir(path):
+    try:
+        os.makedirs(path)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+def init_distributed_mode(args):
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        args["rank"] = int(os.environ["RANK"])
+        args["world_size"] = int(os.environ["WORLD_SIZE"])
+        args["gpu"] = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args["rank"] = int(os.environ["SLURM_PROCID"])
+        args["gpu"] = args["rank"] % torch.cuda.device_count()
+    else:
+        print("Not using distributed mode")
+        args["distributed"] = False
+        return
+
+    args["distributed"] = True
+
+    torch.cuda.set_device(args["gpu"])
+    args["dist_backend"] = "nccl"
+    # print(f"| distributed init (rank {args['rank']}): {args['dist_url']}", flush=True)
+    torch.distributed.init_process_group(
+        backend=args["dist_backend"], init_method=args["dist_url"], world_size=args["world_size"], rank=args["rank"]
+    )
+    torch.distributed.barrier()
+    setup_for_distributed(args["rank"] == 0)
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 65abbb30..2c383b08 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,3 +9,4 @@ matplotlib==3.7.2
 diffusers==0.18.2
 wandb==0.15.7
 lightning==2.1.0rc0
+tensorboard==2.14.0
\ No newline at end of file
diff --git a/tv-detection/README.md b/tv-detection/README.md
new file mode 100644
index 00000000..bb211ba1
--- /dev/null
+++ b/tv-detection/README.md
@@ -0,0 +1,28 @@
+# Object detection reference training scripts
+
+This folder contains reference training scripts for object detection.
+They serve as a log of how to train specific models, to provide baseline
+training and evaluation scripts to quickly bootstrap research.
+
+To execute the example commands below you must install the following:
+
+```
+cython
+pycocotools
+matplotlib
+```
+
+You must also run "prep.py" to download pretrained model weights before 
+launching your training job.
+
+You can then run the training routines for the following models using cli.
+
+### RetinaNet
+```
+isc train ./retinanet_resnet101_fpn.isc
+```
+
+### Mask R-CNN
+```
+isc train ./maskrcnn_resnet101_fpn.isc
+```
\ No newline at end of file
diff --git a/tv-detection/coco_eval.py b/tv-detection/coco_eval.py
new file mode 100644
index 00000000..2a852b39
--- /dev/null
+++ b/tv-detection/coco_eval.py
@@ -0,0 +1,203 @@
+import copy
+import io
+from contextlib import redirect_stdout
+
+import numpy as np
+import pycocotools.mask as mask_util
+import torch
+import utils
+from pycocotools.coco import COCO
+from pycocotools.cocoeval import COCOeval
+
+
+class CocoEvaluator:
+    def __init__(self, coco_gt, iou_types):
+        if not isinstance(iou_types, (list, tuple)):
+            raise TypeError(f"This constructor expects iou_types of type list or tuple, instead  got {type(iou_types)}")
+        coco_gt = copy.deepcopy(coco_gt)
+        self.coco_gt = coco_gt
+
+        self.iou_types = iou_types
+        self.coco_eval = {}
+        for iou_type in iou_types:
+            self.coco_eval[iou_type] = COCOeval(coco_gt, iouType=iou_type)
+
+        self.img_ids = []
+        self.eval_imgs = {k: [] for k in iou_types}
+
+    def update(self, predictions):
+        '''
+        predictions = {
+            image_id: {
+                "boxes": tensor(n x 4), "labels": tensor(n,), "scores": tensor(n,), "masks": tensor(n x 1 x H x W)
+                }, # for n < N detections
+            ... for image_id in batch
+        }
+        '''
+        img_ids = list(np.unique(list(predictions.keys()))) # images seen this batch
+        self.img_ids.extend(img_ids) # catalogue of all images seen
+
+        for iou_type in self.iou_types: # either ["bbox"] or ["bbox", "segm"] or ["bbox", "keypoints"]
+            results = self.prepare(predictions, iou_type)
+
+            with redirect_stdout(io.StringIO()):
+                coco_dt = COCO.loadRes(self.coco_gt, results) if results else COCO()
+            coco_eval = self.coco_eval[iou_type]
+
+            coco_eval.cocoDt = coco_dt
+            coco_eval.params.imgIds = list(img_ids)
+            img_ids, eval_imgs = evaluate(coco_eval)
+
+            self.eval_imgs[iou_type].append(eval_imgs)
+
+    def synchronize_between_processes(self):
+        for iou_type in self.iou_types:
+            self.eval_imgs[iou_type] = np.concatenate(self.eval_imgs[iou_type], 2)
+            create_common_coco_eval(self.coco_eval[iou_type], self.img_ids, self.eval_imgs[iou_type])
+
+    def accumulate(self):
+        for coco_eval in self.coco_eval.values():
+            coco_eval.accumulate()
+
+    def summarize(self):
+        results = []
+        for iou_type, coco_eval in self.coco_eval.items():
+            print(f"IoU metric: {iou_type}")
+            coco_eval.summarize()
+            results += list(coco_eval.stats)
+        return results
+
+    def prepare(self, predictions, iou_type):
+        if iou_type == "bbox":
+            return self.prepare_for_coco_detection(predictions)
+        if iou_type == "segm":
+            return self.prepare_for_coco_segmentation(predictions)
+        if iou_type == "keypoints":
+            return self.prepare_for_coco_keypoint(predictions)
+        raise ValueError(f"Unknown iou type {iou_type}")
+
+    def prepare_for_coco_detection(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "bbox": box,
+                        "score": scores[k],
+                    }
+                    for k, box in enumerate(boxes)
+                ]
+            )
+        return coco_results
+
+    def prepare_for_coco_segmentation(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            scores = prediction["scores"]
+            labels = prediction["labels"]
+            masks = prediction["masks"]
+
+            masks = masks > 0.5
+
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+
+            rles = [
+                mask_util.encode(np.array(mask[0, :, :, np.newaxis], dtype=np.uint8, order="F"))[0] for mask in masks
+            ]
+            for rle in rles:
+                rle["counts"] = rle["counts"].decode("utf-8")
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "segmentation": rle,
+                        "score": scores[k],
+                    }
+                    for k, rle in enumerate(rles)
+                ]
+            )
+        return coco_results
+
+    def prepare_for_coco_keypoint(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+            keypoints = prediction["keypoints"]
+            keypoints = keypoints.flatten(start_dim=1).tolist()
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "keypoints": keypoint,
+                        "score": scores[k],
+                    }
+                    for k, keypoint in enumerate(keypoints)
+                ]
+            )
+        return coco_results
+
+
+def convert_to_xywh(boxes):
+    xmin, ymin, xmax, ymax = boxes.unbind(1)
+    return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=1)
+
+
+def merge(img_ids, eval_imgs):
+    all_img_ids = utils.all_gather(img_ids)
+    all_eval_imgs = utils.all_gather(eval_imgs)
+
+    merged_img_ids = []
+    for p in all_img_ids:
+        merged_img_ids.extend(p)
+
+    merged_eval_imgs = []
+    for p in all_eval_imgs:
+        merged_eval_imgs.append(p)
+
+    merged_img_ids = np.array(merged_img_ids)
+    merged_eval_imgs = np.concatenate(merged_eval_imgs, 2)
+
+    # keep only unique (and in sorted order) images
+    merged_img_ids, idx = np.unique(merged_img_ids, return_index=True)
+    merged_eval_imgs = merged_eval_imgs[..., idx]
+
+    return merged_img_ids, merged_eval_imgs
+
+
+def create_common_coco_eval(coco_eval, img_ids, eval_imgs):
+    img_ids, eval_imgs = merge(img_ids, eval_imgs)
+    img_ids = list(img_ids)
+    eval_imgs = list(eval_imgs.flatten())
+
+    coco_eval.evalImgs = eval_imgs
+    coco_eval.params.imgIds = img_ids
+    coco_eval._paramsEval = copy.deepcopy(coco_eval.params)
+
+
+def evaluate(imgs):
+    with redirect_stdout(io.StringIO()):
+        imgs.evaluate()
+    return imgs.params.imgIds, np.asarray(imgs.evalImgs).reshape(-1, len(imgs.params.areaRng), len(imgs.params.imgIds))
diff --git a/tv-detection/coco_utils.py b/tv-detection/coco_utils.py
new file mode 100644
index 00000000..f40dcdff
--- /dev/null
+++ b/tv-detection/coco_utils.py
@@ -0,0 +1,234 @@
+import os
+
+import torch
+import torch.utils.data
+import torchvision
+import transforms as T
+from pycocotools import mask as coco_mask
+from pycocotools.coco import COCO
+
+
+def convert_coco_poly_to_mask(segmentations, height, width):
+    masks = []
+    for polygons in segmentations:
+        rles = coco_mask.frPyObjects(polygons, height, width)
+        mask = coco_mask.decode(rles)
+        if len(mask.shape) < 3:
+            mask = mask[..., None]
+        mask = torch.as_tensor(mask, dtype=torch.uint8)
+        mask = mask.any(dim=2)
+        masks.append(mask)
+    if masks:
+        masks = torch.stack(masks, dim=0)
+    else:
+        masks = torch.zeros((0, height, width), dtype=torch.uint8)
+    return masks
+
+
+class ConvertCocoPolysToMask:
+    def __call__(self, image, target):
+        w, h = image.size
+
+        image_id = target["image_id"]
+
+        anno = target["annotations"]
+
+        anno = [obj for obj in anno if obj["iscrowd"] == 0]
+
+        boxes = [obj["bbox"] for obj in anno]
+        # guard against no boxes via resizing
+        boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
+        boxes[:, 2:] += boxes[:, :2]
+        boxes[:, 0::2].clamp_(min=0, max=w)
+        boxes[:, 1::2].clamp_(min=0, max=h)
+
+        classes = [obj["category_id"] for obj in anno]
+        classes = torch.tensor(classes, dtype=torch.int64)
+
+        segmentations = [obj["segmentation"] for obj in anno]
+        masks = convert_coco_poly_to_mask(segmentations, h, w)
+
+        keypoints = None
+        if anno and "keypoints" in anno[0]:
+            keypoints = [obj["keypoints"] for obj in anno]
+            keypoints = torch.as_tensor(keypoints, dtype=torch.float32)
+            num_keypoints = keypoints.shape[0]
+            if num_keypoints:
+                keypoints = keypoints.view(num_keypoints, -1, 3)
+
+        keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
+        boxes = boxes[keep]
+        classes = classes[keep]
+        masks = masks[keep]
+        if keypoints is not None:
+            keypoints = keypoints[keep]
+
+        target = {}
+        target["boxes"] = boxes
+        target["labels"] = classes
+        target["masks"] = masks
+        target["image_id"] = image_id
+        if keypoints is not None:
+            target["keypoints"] = keypoints
+
+        # for conversion to coco api
+        area = torch.tensor([obj["area"] for obj in anno])
+        iscrowd = torch.tensor([obj["iscrowd"] for obj in anno])
+        target["area"] = area
+        target["iscrowd"] = iscrowd
+
+        return image, target
+
+
+def _coco_remove_images_without_annotations(dataset, cat_list=None):
+    def _has_only_empty_bbox(anno):
+        return all(any(o <= 1 for o in obj["bbox"][2:]) for obj in anno)
+
+    def _count_visible_keypoints(anno):
+        return sum(sum(1 for v in ann["keypoints"][2::3] if v > 0) for ann in anno)
+
+    min_keypoints_per_image = 10
+
+    def _has_valid_annotation(anno):
+        # if it's empty, there is no annotation
+        if len(anno) == 0:
+            return False
+        # if all boxes have close to zero area, there is no annotation
+        if _has_only_empty_bbox(anno):
+            return False
+        # keypoints task have a slight different criteria for considering
+        # if an annotation is valid
+        if "keypoints" not in anno[0]:
+            return True
+        # for keypoint detection tasks, only consider valid images those
+        # containing at least min_keypoints_per_image
+        if _count_visible_keypoints(anno) >= min_keypoints_per_image:
+            return True
+        return False
+
+    ids = []
+    for ds_idx, img_id in enumerate(dataset.ids):
+        ann_ids = dataset.coco.getAnnIds(imgIds=img_id, iscrowd=None)
+        anno = dataset.coco.loadAnns(ann_ids)
+        if cat_list:
+            anno = [obj for obj in anno if obj["category_id"] in cat_list]
+        if _has_valid_annotation(anno):
+            ids.append(ds_idx)
+
+    dataset = torch.utils.data.Subset(dataset, ids)
+    return dataset
+
+
+def convert_to_coco_api(ds):
+    coco_ds = COCO()
+    # annotation IDs need to start at 1, not 0, see torchvision issue #1530
+    ann_id = 1
+    dataset = {"images": [], "categories": [], "annotations": []}
+    categories = set()
+    for img_idx in range(len(ds)):
+        # find better way to get target
+        # targets = ds.get_annotations(img_idx)
+        img, targets = ds[img_idx]
+        image_id = targets["image_id"]
+        img_dict = {}
+        img_dict["id"] = image_id
+        img_dict["height"] = img.shape[-2]
+        img_dict["width"] = img.shape[-1]
+        dataset["images"].append(img_dict)
+        bboxes = targets["boxes"].clone()
+        bboxes[:, 2:] -= bboxes[:, :2]
+        bboxes = bboxes.tolist()
+        labels = targets["labels"].tolist()
+        areas = targets["area"].tolist()
+        iscrowd = targets["iscrowd"].tolist()
+        if "masks" in targets:
+            masks = targets["masks"]
+            # make masks Fortran contiguous for coco_mask
+            masks = masks.permute(0, 2, 1).contiguous().permute(0, 2, 1)
+        if "keypoints" in targets:
+            keypoints = targets["keypoints"]
+            keypoints = keypoints.reshape(keypoints.shape[0], -1).tolist()
+        num_objs = len(bboxes)
+        for i in range(num_objs):
+            ann = {}
+            ann["image_id"] = image_id
+            ann["bbox"] = bboxes[i]
+            ann["category_id"] = labels[i]
+            categories.add(labels[i])
+            ann["area"] = areas[i]
+            ann["iscrowd"] = iscrowd[i]
+            ann["id"] = ann_id
+            if "masks" in targets:
+                ann["segmentation"] = coco_mask.encode(masks[i].numpy())
+            if "keypoints" in targets:
+                ann["keypoints"] = keypoints[i]
+                ann["num_keypoints"] = sum(k != 0 for k in keypoints[i][2::3])
+            dataset["annotations"].append(ann)
+            ann_id += 1
+    dataset["categories"] = [{"id": i} for i in sorted(categories)]
+    coco_ds.dataset = dataset
+    coco_ds.createIndex()
+    return coco_ds
+
+
+def get_coco_api_from_dataset(dataset):
+    # FIXME: This is... awful?
+    for _ in range(10):
+        if isinstance(dataset, torchvision.datasets.CocoDetection):
+            break
+        if isinstance(dataset, torch.utils.data.Subset):
+            dataset = dataset.dataset
+    if isinstance(dataset, torchvision.datasets.CocoDetection):
+        return dataset.coco
+    return convert_to_coco_api(dataset)
+
+
+class CocoDetection(torchvision.datasets.CocoDetection):
+    def __init__(self, img_folder, ann_file, transforms):
+        super().__init__(img_folder, ann_file)
+        self._transforms = transforms
+
+    def __getitem__(self, idx):
+        img, target = super().__getitem__(idx)
+        image_id = self.ids[idx]
+        target = dict(image_id=image_id, annotations=target)
+        if self._transforms is not None:
+            img, target = self._transforms(img, target)
+        return img, target
+
+
+def get_coco(root, image_set, transforms, mode="instances", use_v2=False, with_masks=False):
+    anno_file_template = "{}_{}2017.json"
+    PATHS = {
+        "train": ("train2017", os.path.join("annotations", anno_file_template.format(mode, "train"))),
+        "val": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val"))),
+        # "train": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val")))
+    }
+
+    img_folder, ann_file = PATHS[image_set]
+    img_folder = os.path.join(root, img_folder)
+    ann_file = os.path.join(root, ann_file)
+
+    if use_v2:
+        from torchvision.datasets import wrap_dataset_for_transforms_v2
+
+        dataset = torchvision.datasets.CocoDetection(img_folder, ann_file, transforms=transforms)
+        target_keys = ["boxes", "labels", "image_id"]
+        if with_masks:
+            target_keys += ["masks"]
+        dataset = wrap_dataset_for_transforms_v2(dataset, target_keys=target_keys)
+    else:
+        # TODO: handle with_masks for V1?
+        t = [ConvertCocoPolysToMask()]
+        if transforms is not None:
+            t.append(transforms)
+        transforms = T.Compose(t)
+
+        dataset = CocoDetection(img_folder, ann_file, transforms=transforms)
+
+    if image_set == "train":
+        dataset = _coco_remove_images_without_annotations(dataset)
+
+    # dataset = torch.utils.data.Subset(dataset, [i for i in range(500)])
+
+    return dataset
diff --git a/tv-detection/engine.py b/tv-detection/engine.py
new file mode 100644
index 00000000..ff97a99a
--- /dev/null
+++ b/tv-detection/engine.py
@@ -0,0 +1,259 @@
+import math
+import sys
+# from itertools import product
+
+import torch
+import torchvision.models.detection.mask_rcnn
+import torch.distributed as dist
+import utils
+from coco_eval import CocoEvaluator
+from coco_utils import get_coco_api_from_dataset
+from cycling_utils import atomic_torch_save
+
+from torch.utils.tensorboard import SummaryWriter
+
+def train_one_epoch(
+        model, optimizer, data_loader_train, train_sampler, test_sampler,
+        lr_scheduler, warmup_lr_scheduler, args, device, coco_evaluator,
+        epoch, scaler=None, timer=None, metrics=None,
+    ):
+
+    model.train()
+
+    timer.report('training preliminaries')
+
+    print(f'\nTraining / resuming epoch {epoch} from training step {train_sampler.progress}\n')
+
+    for i, (images, targets) in enumerate(data_loader_train):
+
+        images = list(image.to(device) for image in images)
+        targets = [{k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in t.items()} for t in targets]
+        timer.report(f'Epoch: {epoch} batch {train_sampler.progress}: moving batch data to device')
+        # print(f"First 2 image shapes: {images[0].shape}, {images[1].shape}")
+
+        optimizer.zero_grad()
+
+        with torch.cuda.amp.autocast(enabled=scaler is not None):
+            assert len(targets) > 0, "Targets iterable of length 0, will return infinite loss."
+            loss_dict = model(images, targets)
+
+            # CHECK IF NUMERIC ERROR HAS OCCURRED AND IF SO, SKIP THIS BATCH
+            check_0 = 1 if torch.tensor([torch.isnan(v) for v in loss_dict.values()]).any() else 0
+            check_1 = 1 if not all([math.isfinite(v) for v in loss_dict.values()]) else 0
+            check_tensor = torch.tensor([check_0, check_1], requires_grad=False, device=device)
+            dist.all_reduce(check_tensor, op=dist.ReduceOp.SUM)
+            
+            if check_tensor.sum() > 0:
+                print(f"CONTINUE CONDITION - NaN: {check_tensor[0].item()}, Infinite: {check_tensor[1].item()}")
+
+                # reset optimizer to prevent momentum carrying model into same issue
+                del optimizer, images, targets
+                torch.cuda.empty_cache()
+
+                if args.norm_weight_decay is None:
+                    parameters = [p for p in model.parameters() if p.requires_grad]
+                else:
+                    param_groups = torchvision.ops._utils.split_normalization_params(model)
+                    wd_groups = [args.norm_weight_decay, args.weight_decay]
+                    parameters = [{"params": p, "weight_decay": w} for p, w in zip(param_groups, wd_groups) if p]
+
+                opt_name = args.opt.lower()
+                if opt_name.startswith("sgd"):
+                    optimizer = torch.optim.SGD(
+                        parameters, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay,
+                        nesterov="nesterov" in opt_name,
+                    )
+                elif opt_name == "adamw":
+                    optimizer = torch.optim.AdamW(parameters, lr=args.lr, weight_decay=args.weight_decay)
+
+                # Advance sampler to try next batch
+                train_sampler.advance()
+                continue
+
+            losses = sum(loss for loss in loss_dict.values())
+            timer.report(f'Epoch: {epoch} batch {train_sampler.progress}: forward pass')
+
+        # trying gradient clipping to prevent gradient issues with retinanet...
+        if args.model == 'retinanet_resnet101_fpn':
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+            timer.report(f'Epoch: {epoch} batch {train_sampler.progress}: graient clipping')
+
+        if scaler is not None:
+            scaler.scale(losses).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            losses.backward()
+            optimizer.step()
+        warmup_lr_scheduler.step()
+
+        timer.report(f'Epoch: {epoch} batch {train_sampler.progress}: backward pass')
+
+        # reduce losses over all GPUs for logging purposes
+        loss_dict_reduced = utils.reduce_dict(loss_dict)
+        losses_reduced = sum(loss for loss in loss_dict_reduced.values())
+        loss_value = losses_reduced.item()
+        timer.report(f'Epoch: {epoch} batch {train_sampler.progress}: computing loss')
+
+        if not math.isfinite(loss_value):
+            print(f"Loss is {loss_value}, stopping training")
+            print(loss_dict_reduced)
+            sys.exit(1)
+
+        metrics["train"].update({"images_seen": len(images) ,"loss": loss_value})
+        metrics["train"].update({k:v.item() for k,v in loss_dict_reduced.items()})
+        metrics["train"].reduce() # Gather results from all nodes
+        
+        report_metrics = [m for m in metrics["train"].local if m != "images_seen"]
+        images_seen = metrics["train"].local["images_seen"]
+        vals = [metrics["train"].local[m]/images_seen for m in report_metrics]
+        rpt = ", ".join([f"{m}: {v:,.3f}" for m,v in zip(report_metrics, vals)])
+        print(f"EPOCH: [{epoch}], BATCH: [{train_sampler.progress}/{len(train_sampler)}], "+rpt)
+
+        metrics["train"].reset_local()
+
+        print(f"Saving checkpoint at epoch {epoch} train batch {train_sampler.progress}")
+        train_sampler.advance()
+
+        if train_sampler.progress == len(train_sampler):
+            metrics["train"].end_epoch()
+
+        if utils.is_main_process() and train_sampler.progress % 1 == 0: # Checkpointing every batch
+
+            writer = SummaryWriter(log_dir=args.tboard_path)
+            for metric,val in zip(report_metrics, vals):
+                writer.add_scalar("Train/"+metric, val, train_sampler.progress + epoch * len(train_sampler))
+            writer.flush()
+            writer.close()
+                
+            checkpoint = {
+                "args": args,
+                "epoch": epoch,
+                "model": model.module.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "warmup_lr_scheduler": warmup_lr_scheduler.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "test_sampler": test_sampler.state_dict(),
+                # Evaluator state variables
+                "img_ids": coco_evaluator.img_ids, # catalogue of images seen already
+                "eval_imgs": coco_evaluator.eval_imgs, # image evaluations
+                "metrics": metrics,
+            }
+            if args.amp:
+                checkpoint["scaler"] = scaler.state_dict()
+            timer = atomic_torch_save(checkpoint, args.resume, timer)
+    
+    lr_scheduler.step() # OUTER LR_SCHEDULER STEP EACH EPOCH
+    return model, timer, metrics
+
+def _get_iou_types(model):
+    model_without_ddp = model
+    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+        model_without_ddp = model.module
+    iou_types = ["bbox"]
+    if isinstance(model_without_ddp, torchvision.models.detection.MaskRCNN):
+        iou_types.append("segm")
+    if isinstance(model_without_ddp, torchvision.models.detection.KeypointRCNN):
+        iou_types.append("keypoints")
+    return iou_types
+
+@torch.inference_mode()
+def evaluate(
+    model, data_loader_test, epoch, test_sampler, args, coco_evaluator,
+    optimizer, lr_scheduler, warmup_lr_scheduler, train_sampler,
+    device, scaler=None, timer=None, metrics=None,
+):
+
+    timer.report('starting evaluation routine')
+
+    n_threads = torch.get_num_threads()
+    torch.set_num_threads(1)
+    cpu_device = torch.device("cpu")
+    model.eval()
+
+    timer.report('evaluation preliminaries')
+
+    test_step = test_sampler.progress // data_loader_test.batch_size
+    print(f'\nEvaluating / resuming epoch {epoch} from eval step {test_step}\n')
+
+    for i, (images, targets) in enumerate(data_loader_test):
+
+        images = list(img.to(device) for img in images)
+        timer.report(f'Epoch {epoch} batch: {test_step} moving to device')
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+        outputs = model(images)
+        timer.report(f'Epoch {epoch} batch: {test_step} forward through model')
+
+        outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
+        timer.report(f'Epoch {epoch} batch: {test_step} outputs back to cpu')
+
+        res = {target["image_id"]: output for target, output in zip(targets, outputs)}
+        # res = {img_id: {'boxes': T, 'labels': T, 'scores': T, 'masks': T}, ...}
+        coco_evaluator.update(res)
+        timer.report(f'Epoch {epoch} batch: {test_step} update evaluator')
+
+        print(f"Saving checkpoint at epoch {epoch} eval batch {test_step}")
+        test_sampler.advance(len(images))
+        test_step = test_sampler.progress // data_loader_test.batch_size
+
+        if utils.is_main_process() and test_step % 1 == 0: # Checkpointing every batch
+            checkpoint = {
+                "args": args,
+                "epoch": epoch,
+                "model": model.module.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "warmup_lr_scheduler": warmup_lr_scheduler.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "test_sampler": test_sampler.state_dict(),
+                # Evaluator state variables
+                "img_ids": coco_evaluator.img_ids, # catalogue of images seen already
+                "eval_imgs": coco_evaluator.eval_imgs, # image evaluations
+                "metrics": metrics,
+            }
+            if args.amp:
+                checkpoint["scaler"] = scaler.state_dict()
+            timer = atomic_torch_save(checkpoint, args.resume, timer)
+    
+    # gather the stats from all processes
+    coco_evaluator.synchronize_between_processes()
+
+    # accumulate predictions from all images
+    coco_evaluator.accumulate()
+    results = coco_evaluator.summarize()
+
+    metric_names = [
+        "bbox/AP", "bbox/AP-50", "bbox/AP-75", "bbox/AP-S", "bbox/AP-M", "bbox/AP-L", 
+        "bbox/AR-MD1", "bbox/AR-MD10", "bbox/AR-MD100", "bbox/AR-S", "bbox/AR-M", "bbox/AR-L",
+        "segm/AP", "segm/AP-50", "segm/AP-75", "segm/AP-S", "segm/AP-M", "segm/AP-L", 
+        "segm/AR-MD1", "segm/AR-MD10", "segm/AR-MD100", "segm/AR-S", "segm/AR-M", "segm/AR-L"
+    ]
+
+    metrics["val"].update({name: val for name,val in zip(metric_names, results)})
+    metrics["val"].reduce()
+    # Normalise validation metrics by world_size
+    ngpus = dist.get_world_size()
+    metrics["val"].agg = {k:v/ngpus for k,v in metrics["val"].agg.items()}
+    metrics["val"].end_epoch()
+
+    if utils.is_main_process():
+        writer = SummaryWriter(log_dir=args.tboard_path)
+        for name,val in metrics["val"].epoch_reports[-1].items():
+            writer.add_scalar("Val/"+name, val/ngpus, epoch)
+        writer.flush()
+        writer.close()
+
+    torch.set_num_threads(n_threads)
+
+    # Reset the coco evaluator at the end of the epoch
+    coco = get_coco_api_from_dataset(data_loader_test.dataset)
+    iou_types = _get_iou_types(model)
+    coco_evaluator = CocoEvaluator(coco, iou_types)
+
+    timer.report('evaluator accumulation, summarization, and reset')
+
+    return coco_evaluator, timer, metrics
diff --git a/tv-detection/engine_SS.py b/tv-detection/engine_SS.py
new file mode 100644
index 00000000..9a43586c
--- /dev/null
+++ b/tv-detection/engine_SS.py
@@ -0,0 +1,257 @@
+import math
+import sys
+from itertools import product
+
+import torch
+import torchvision.models.detection.mask_rcnn
+import torch.distributed as dist 
+import utils
+from coco_eval import CocoEvaluator
+from coco_utils import get_coco_api_from_dataset
+from cycling_utils import atomic_torch_save
+
+from torch.utils.tensorboard import SummaryWriter
+
+def train_one_epoch(
+        model, optimizer, data_loader_train, train_sampler, test_sampler,
+        lr_scheduler, warmup_lr_scheduler, args, device, coco_evaluator,
+        epoch, scaler=None, timer=None, metrics=None,
+    ):
+
+    model.train()
+
+    timer.report('training preliminaries')
+
+    print(f'\nTraining / resuming epoch {epoch} from training step {train_sampler.progress}\n')
+
+    for images, targets in data_loader_train:
+
+        images = list(image.to(device) for image in images)
+        targets = [{k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in t.items()} for t in targets]
+        timer.report(f'Epoch: {epoch} batch {train_sampler.progress}: moving batch data to device')
+        # print(f"First 2 image shapes: {images[0].shape}, {images[1].shape}")
+
+        optimizer.zero_grad()
+
+        with torch.cuda.amp.autocast(enabled=scaler is not None):
+            assert len(targets) > 0, "Targets iterable of length 0, will return infinite loss."
+            loss_dict = model(images, targets)
+
+            # CHECK IF NUMERIC ERROR HAS OCCURRED AND IF SO, SKIP THIS BATCH
+            check_0 = 1 if torch.tensor([torch.isnan(v) for v in loss_dict.values()]).any() else 0
+            check_1 = 1 if not all([math.isfinite(v) for v in loss_dict.values()]) else 0
+            check_tensor = torch.tensor([check_0, check_1], requires_grad=False, device=device)
+            dist.all_reduce(check_tensor, op=dist.ReduceOp.SUM)
+            if check_tensor.sum() > 0:
+                print(f"CONTINUE CONDITION - NaN: {check_tensor[0].item()}, Infinite: {check_tensor[1].item()}")
+
+                # reset optimizer to prevent momentum carrying model into same issue
+                del optimizer
+                if args.norm_weight_decay is None:
+                    parameters = [p for p in model.parameters() if p.requires_grad]
+                else:
+                    param_groups = torchvision.ops._utils.split_normalization_params(model)
+                    wd_groups = [args.norm_weight_decay, args.weight_decay]
+                    parameters = [{"params": p, "weight_decay": w} for p, w in zip(param_groups, wd_groups) if p]
+
+                opt_name = args.opt.lower()
+                if opt_name.startswith("sgd"):
+                    optimizer = torch.optim.SGD(
+                        parameters, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay,
+                        nesterov="nesterov" in opt_name,
+                    )
+                elif opt_name == "adamw":
+                    optimizer = torch.optim.AdamW(parameters, lr=args.lr, weight_decay=args.weight_decay)
+
+                # Advance sampler to try next batch
+                train_sampler.advance()
+                continue
+
+            losses = sum(loss for loss in loss_dict.values())
+            timer.report(f'Epoch: {epoch} batch {train_sampler.progress}: forward pass')
+
+        # trying gradient clipping to prevent gradient issues with retinanet...
+        if args.model == 'retinanet_resnet101_fpn':
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+            timer.report(f'Epoch: {epoch} batch {train_sampler.progress}: graient clipping')
+
+        if scaler is not None:
+            scaler.scale(losses).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            losses.backward()
+            optimizer.step()
+        warmup_lr_scheduler.step()
+
+        timer.report(f'Epoch: {epoch} batch {train_sampler.progress}: backward pass')
+
+        # reduce losses over all GPUs for logging purposes
+        loss_dict_reduced = utils.reduce_dict(loss_dict)
+        losses_reduced = sum(loss for loss in loss_dict_reduced.values())
+        loss_value = losses_reduced.item()
+        timer.report(f'Epoch: {epoch} batch {train_sampler.progress}: computing loss')
+
+        if not math.isfinite(loss_value):
+            print(f"Loss is {loss_value}, stopping training")
+            print(loss_dict_reduced)
+            sys.exit(1)
+
+        metrics["train"].update({"images_seen": len(images) ,"loss": loss_value})
+        metrics["train"].update({k:v.item() for k,v in loss_dict_reduced.items()})
+        metrics["train"].reduce() # Gather results from all nodes
+        
+        report_metrics = [m for m in metrics["train"].local if m != "images_seen"]
+        images_seen = metrics["train"].local["images_seen"]
+        vals = [metrics["train"].local[m]/images_seen for m in report_metrics]
+        rpt = ", ".join([f"{m}: {v:,.3f}" for m,v in zip(report_metrics, vals)])
+        print(f"EPOCH: [{epoch}], BATCH: [{train_sampler.progress}/{len(train_sampler)}], "+rpt)
+
+        metrics["train"].reset_local()
+
+        print(f"Saving checkpoint at epoch {epoch} train batch {train_sampler.progress}")
+        train_sampler.advance()
+
+        if train_sampler.progress == len(train_sampler):
+            metrics["train"].end_epoch()
+
+        if utils.is_main_process() and train_sampler.progress % 1 == 0: # Checkpointing every batch
+
+            writer = SummaryWriter(log_dir=args.tboard_path)
+            for metric,val in zip(report_metrics, vals):
+                writer.add_scalar("Train/"+metric, val, train_sampler.progress + epoch * len(train_sampler))
+            writer.flush()
+            writer.close()
+                
+            checkpoint = {
+                "args": args,
+                "epoch": epoch,
+                "model": model.module.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "warmup_lr_scheduler": warmup_lr_scheduler.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "test_sampler": test_sampler.state_dict(),
+                # Evaluator state variables
+                "img_ids": coco_evaluator.img_ids, # catalogue of images seen already
+                "eval_imgs": coco_evaluator.eval_imgs, # image evaluations
+                "metrics": metrics,
+            }
+            if args.amp:
+                checkpoint["scaler"] = scaler.state_dict()
+            timer = atomic_torch_save(checkpoint, args.resume, timer)
+
+    lr_scheduler.step() # OUTER LR_SCHEDULER STEP EACH EPOCH
+    return model, timer, metrics
+
+def _get_iou_types(model):
+    model_without_ddp = model
+    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+        model_without_ddp = model.module
+    iou_types = ["bbox"]
+    if isinstance(model_without_ddp, torchvision.models.detection.MaskRCNN):
+        iou_types.append("segm")
+    if isinstance(model_without_ddp, torchvision.models.detection.KeypointRCNN):
+        iou_types.append("keypoints")
+    return iou_types
+
+@torch.inference_mode()
+def evaluate(
+    model, data_loader_test, epoch, test_sampler, args, coco_evaluator,
+    optimizer, lr_scheduler, warmup_lr_scheduler, train_sampler,
+    device, scaler=None, timer=None, metrics=None,
+):
+
+    timer.report('starting evaluation routine')
+
+    n_threads = torch.get_num_threads()
+    torch.set_num_threads(1)
+    cpu_device = torch.device("cpu")
+    model.eval()
+
+    timer.report('evaluation preliminaries')
+
+    test_step = test_sampler.progress // data_loader_test.batch_size
+    print(f'\nEvaluating / resuming epoch {epoch} from eval step {test_step}\n')
+
+    for images, targets in data_loader_test:
+
+        images = list(img.to(device) for img in images)
+        timer.report(f'Epoch {epoch} batch: {test_step} moving to device')
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+        outputs = model(images)
+        timer.report(f'Epoch {epoch} batch: {test_step} forward through model')
+
+        outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
+        timer.report(f'Epoch {epoch} batch: {test_step} outputs back to cpu')
+
+        res = {target["image_id"]: output for target, output in zip(targets, outputs)}
+        # res = {img_id: {'boxes': T, 'labels': T, 'scores': T, 'masks': T}, ...}
+        coco_evaluator.update(res)
+        timer.report(f'Epoch {epoch} batch: {test_step} update evaluator')
+
+        print(f"Saving checkpoint at epoch {epoch} eval batch {test_step}")
+        test_sampler.advance(len(images))
+        test_step = test_sampler.progress // data_loader_test.batch_size
+
+        if utils.is_main_process() and test_step % 1 == 0: # Checkpointing every batch
+            checkpoint = {
+                "args": args,
+                "epoch": epoch,
+                "model": model.module.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "warmup_lr_scheduler": warmup_lr_scheduler.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "test_sampler": test_sampler.state_dict(),
+                # Evaluator state variables
+                "img_ids": coco_evaluator.img_ids, # catalogue of images seen already
+                "eval_imgs": coco_evaluator.eval_imgs, # image evaluations
+                "metrics": metrics,
+            }
+            if args.amp:
+                checkpoint["scaler"] = scaler.state_dict()
+            timer = atomic_torch_save(checkpoint, args.resume, timer)
+
+    # gather the stats from all processes
+    coco_evaluator.synchronize_between_processes()
+
+    # accumulate predictions from all images
+    coco_evaluator.accumulate()
+    results = coco_evaluator.summarize()
+
+    # metric_A = ["bbox-", "segm-"]
+    # metric_B = ["AP", "AR"]
+    # metric_C = ["", "50", "75", "-S", "-M", "-L"]
+    # metric_names = ["".join(t) for t in product(metric_A, metric_B, metric_C)]
+    metric_names = [
+        "bbox/AP", "bbox/AP-50", "bbox/AP-75", "bbox/AP-S", "bbox/AP-M", "bbox/AP-L", 
+        "bbox/AR-MD1", "bbox/AR-MD10", "bbox/AR-MD100", "bbox/AR-S", "bbox/AR-M", "bbox/AR-L"
+    ] + [
+        "segm/AP", "segm/AP-50", "segm/AP-75", "segm/AP-S", "segm/AP-M", "segm/AP-L", 
+        "segm/AR-MD1", "segm/AR-MD10", "segm/AR-MD100", "segm/AR-S", "segm/AR-M", "segm/AR-L"
+    ]
+    metrics["val"].update({name: val for name,val in zip(metric_names, results)})
+    metrics["val"].reduce()
+    metrics["val"].end_epoch()
+
+    if utils.is_main_process():
+        writer = SummaryWriter(log_dir=args.tboard_path)
+        for name,val in metrics["val"].epoch_reports[-1].items():
+            writer.add_scalar("Val/"+name, val, epoch)
+        writer.flush()
+        writer.close()
+
+    torch.set_num_threads(n_threads)
+
+    # Reset the coco evaluator at the end of the epoch
+    coco = get_coco_api_from_dataset(data_loader_test.dataset)
+    iou_types = _get_iou_types(model)
+    coco_evaluator = CocoEvaluator(coco, iou_types)
+
+    timer.report('evaluator accumulation, summarization, and reset')
+
+    return coco_evaluator, timer, metrics
diff --git a/tv-detection/group_by_aspect_ratio.py b/tv-detection/group_by_aspect_ratio.py
new file mode 100644
index 00000000..7ed30b0b
--- /dev/null
+++ b/tv-detection/group_by_aspect_ratio.py
@@ -0,0 +1,196 @@
+import bisect
+import copy
+import math
+from collections import defaultdict
+from itertools import chain, repeat
+
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from PIL import Image
+from torch.utils.data.sampler import BatchSampler, Sampler
+from torch.utils.model_zoo import tqdm
+
+
+def _repeat_to_at_least(iterable, n):
+    repeat_times = math.ceil(n / len(iterable))
+    repeated = chain.from_iterable(repeat(iterable, repeat_times))
+    return list(repeated)
+
+
+class GroupedBatchSampler(BatchSampler):
+    """
+    Wraps another sampler to yield a mini-batch of indices.
+    It enforces that the batch only contain elements from the same group.
+    It also tries to provide mini-batches which follows an ordering which is
+    as close as possible to the ordering from the original sampler.
+    Args:
+        sampler (Sampler): Base sampler.
+        group_ids (list[int]): If the sampler produces indices in range [0, N),
+            `group_ids` must be a list of `N` ints which contains the group id of each sample.
+            The group ids must be a continuous set of integers starting from
+            0, i.e. they must be in the range [0, num_groups).
+        batch_size (int): Size of mini-batch.
+    """
+
+    def __init__(self, sampler, group_ids, batch_size):
+        if not isinstance(sampler, Sampler):
+            raise ValueError(f"sampler should be an instance of torch.utils.data.Sampler, but got sampler={sampler}")
+        self.sampler = sampler
+        self.group_ids = group_ids
+        self.batch_size = batch_size
+
+    def __iter__(self):
+        buffer_per_group = defaultdict(list)
+        samples_per_group = defaultdict(list)
+
+        num_batches = 0
+        for idx in self.sampler:
+            group_id = self.group_ids[idx]
+            buffer_per_group[group_id].append(idx)
+            samples_per_group[group_id].append(idx)
+            if len(buffer_per_group[group_id]) == self.batch_size:
+                yield buffer_per_group[group_id]
+                num_batches += 1
+                del buffer_per_group[group_id]
+            assert len(buffer_per_group[group_id]) < self.batch_size
+
+        # now we have run out of elements that satisfy
+        # the group criteria, let's return the remaining
+        # elements so that the size of the sampler is
+        # deterministic
+        expected_num_batches = len(self)
+        num_remaining = expected_num_batches - num_batches
+        if num_remaining > 0:
+            # for the remaining batches, take first the buffers with the largest number
+            # of elements
+            for group_id, _ in sorted(buffer_per_group.items(), key=lambda x: len(x[1]), reverse=True):
+                remaining = self.batch_size - len(buffer_per_group[group_id])
+                samples_from_group_id = _repeat_to_at_least(samples_per_group[group_id], remaining)
+                buffer_per_group[group_id].extend(samples_from_group_id[:remaining])
+                assert len(buffer_per_group[group_id]) == self.batch_size
+                yield buffer_per_group[group_id]
+                num_remaining -= 1
+                if num_remaining == 0:
+                    break
+        assert num_remaining == 0
+
+    def __len__(self):
+        return len(self.sampler) // self.batch_size
+
+
+def _compute_aspect_ratios_slow(dataset, indices=None):
+    print(
+        "Your dataset doesn't support the fast path for "
+        "computing the aspect ratios, so will iterate over "
+        "the full dataset and load every image instead. "
+        "This might take some time..."
+    )
+    if indices is None:
+        indices = range(len(dataset))
+
+    class SubsetSampler(Sampler):
+        def __init__(self, indices):
+            self.indices = indices
+
+        def __iter__(self):
+            return iter(self.indices)
+
+        def __len__(self):
+            return len(self.indices)
+
+    sampler = SubsetSampler(indices)
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=1,
+        sampler=sampler,
+        num_workers=14,  # you might want to increase it for faster processing
+        collate_fn=lambda x: x[0],
+    )
+    aspect_ratios = []
+    with tqdm(total=len(dataset)) as pbar:
+        for _i, (img, _) in enumerate(data_loader):
+            pbar.update(1)
+            height, width = img.shape[-2:]
+            aspect_ratio = float(width) / float(height)
+            aspect_ratios.append(aspect_ratio)
+    return aspect_ratios
+
+
+def _compute_aspect_ratios_custom_dataset(dataset, indices=None):
+    if indices is None:
+        indices = range(len(dataset))
+    aspect_ratios = []
+    for i in indices:
+        height, width = dataset.get_height_and_width(i)
+        aspect_ratio = float(width) / float(height)
+        aspect_ratios.append(aspect_ratio)
+    return aspect_ratios
+
+
+def _compute_aspect_ratios_coco_dataset(dataset, indices=None):
+    if indices is None:
+        indices = range(len(dataset))
+    aspect_ratios = []
+    for i in indices:
+        img_info = dataset.coco.imgs[dataset.ids[i]]
+        aspect_ratio = float(img_info["width"]) / float(img_info["height"])
+        aspect_ratios.append(aspect_ratio)
+    return aspect_ratios
+
+
+def _compute_aspect_ratios_voc_dataset(dataset, indices=None):
+    if indices is None:
+        indices = range(len(dataset))
+    aspect_ratios = []
+    for i in indices:
+        # this doesn't load the data into memory, because PIL loads it lazily
+        width, height = Image.open(dataset.images[i]).size
+        aspect_ratio = float(width) / float(height)
+        aspect_ratios.append(aspect_ratio)
+    return aspect_ratios
+
+
+def _compute_aspect_ratios_subset_dataset(dataset, indices=None):
+    if indices is None:
+        indices = range(len(dataset))
+
+    ds_indices = [dataset.indices[i] for i in indices]
+    return compute_aspect_ratios(dataset.dataset, ds_indices)
+
+
+def compute_aspect_ratios(dataset, indices=None):
+    if hasattr(dataset, "get_height_and_width"):
+        return _compute_aspect_ratios_custom_dataset(dataset, indices)
+
+    if isinstance(dataset, torchvision.datasets.CocoDetection):
+        return _compute_aspect_ratios_coco_dataset(dataset, indices)
+
+    if isinstance(dataset, torchvision.datasets.VOCDetection):
+        return _compute_aspect_ratios_voc_dataset(dataset, indices)
+
+    if isinstance(dataset, torch.utils.data.Subset):
+        return _compute_aspect_ratios_subset_dataset(dataset, indices)
+
+    # slow path
+    return _compute_aspect_ratios_slow(dataset, indices)
+
+
+def _quantize(x, bins):
+    bins = copy.deepcopy(bins)
+    bins = sorted(bins)
+    quantized = list(map(lambda y: bisect.bisect_right(bins, y), x))
+    return quantized
+
+
+def create_aspect_ratio_groups(dataset, k=0):
+    aspect_ratios = compute_aspect_ratios(dataset) # list of aspect ratios for each image in the dataset
+    bins = (2 ** np.linspace(-1, 1, 2 * k + 1)).tolist() if k > 0 else [1.0]
+    groups = _quantize(aspect_ratios, bins) # list of bin indexes to which each image belongs
+    # count number of elements per group
+    counts = np.unique(groups, return_counts=True)[1]
+    fbins = [0] + bins + [np.inf]
+    print(f"Using {fbins} as bins for aspect ratio quantization")
+    print(f"Count of instances per bin: {counts}")
+    return groups
diff --git a/tv-detection/maskrcnn_resnet101_fpn.isc b/tv-detection/maskrcnn_resnet101_fpn.isc
new file mode 100644
index 00000000..cb373319
--- /dev/null
+++ b/tv-detection/maskrcnn_resnet101_fpn.isc
@@ -0,0 +1,6 @@
+experiment_name = "maskrcnn_resnet101_fpn"
+gpu_type = "24GB VRAM GPU"
+nnodes = 11
+venv_path = "/mnt/Client/Adamstn3rh22tykvgyhdkclook3rnk7q/adaadam4qalumfvjdstjpx7zyvlebh2u/.venv/bin/activate"
+output_path = "/mnt/Client/Adamstn3rh22tykvgyhdkclook3rnk7q/adaadam4qalumfvjdstjpx7zyvlebh2u/outputs/maskrcnn_resnet101_fpn"
+command = "train_cycling.py --dataset coco --model maskrcnn_resnet101_fpn --epochs 26 --lr-steps 16 22 -b 2 --aspect-ratio-group-factor 3 --data-path=/mnt/.node1/Open-Datasets/coco --resume $OUTPUT_PATH/checkpoint.isc --tboard-path $OUTPUT_PATH/tb"
diff --git a/tv-detection/prep.py b/tv-detection/prep.py
new file mode 100644
index 00000000..d64ec4db
--- /dev/null
+++ b/tv-detection/prep.py
@@ -0,0 +1,3 @@
+# Downloading ResNet101 backbone weights ahead of training
+from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
+_  = resnet_fpn_backbone(backbone_name="resnet101", weights="ResNet101_Weights.IMAGENET1K_V1")
\ No newline at end of file
diff --git a/tv-detection/presets.py b/tv-detection/presets.py
new file mode 100644
index 00000000..e9b6d56c
--- /dev/null
+++ b/tv-detection/presets.py
@@ -0,0 +1,114 @@
+from collections import defaultdict
+
+import torch
+import transforms as reference_transforms
+
+
+def get_modules(use_v2):
+    # We need a protected import to avoid the V2 warning in case just V1 is used
+    if use_v2:
+        import torchvision.transforms.v2
+        import torchvision.tv_tensors
+
+        return torchvision.transforms.v2, torchvision.tv_tensors
+    else:
+        return reference_transforms, None
+
+
+class DetectionPresetTrain:
+    # Note: this transform assumes that the input to forward() are always PIL
+    # images, regardless of the backend parameter.
+    def __init__(
+        self,
+        *,
+        data_augmentation,
+        hflip_prob=0.5,
+        mean=(123.0, 117.0, 104.0),
+        backend="pil",
+        use_v2=False,
+    ):
+
+        T, tv_tensors = get_modules(use_v2)
+
+        transforms = []
+        backend = backend.lower()
+        if backend == "tv_tensor":
+            transforms.append(T.ToImage())
+        elif backend == "tensor":
+            transforms.append(T.PILToTensor())
+        elif backend != "pil":
+            raise ValueError(f"backend can be 'tv_tensor', 'tensor' or 'pil', but got {backend}")
+
+        if data_augmentation == "hflip":
+            transforms += [T.RandomHorizontalFlip(p=hflip_prob)]
+        elif data_augmentation == "lsj":
+            transforms += [
+                T.ScaleJitter(target_size=(1024, 1024), antialias=True),
+                # TODO: FixedSizeCrop below doesn't work on tensors!
+                reference_transforms.FixedSizeCrop(size=(1024, 1024), fill=mean),
+                T.RandomHorizontalFlip(p=hflip_prob),
+            ]
+        elif data_augmentation == "multiscale":
+            transforms += [
+                T.RandomShortestSize(min_size=(480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=1333),
+                T.RandomHorizontalFlip(p=hflip_prob),
+            ]
+        elif data_augmentation == "ssd":
+            fill = defaultdict(lambda: mean, {tv_tensors.Mask: 0}) if use_v2 else list(mean)
+            transforms += [
+                T.RandomPhotometricDistort(),
+                T.RandomZoomOut(fill=fill),
+                T.RandomIoUCrop(),
+                T.RandomHorizontalFlip(p=hflip_prob),
+            ]
+        elif data_augmentation == "ssdlite":
+            transforms += [
+                T.RandomIoUCrop(),
+                T.RandomHorizontalFlip(p=hflip_prob),
+            ]
+        else:
+            raise ValueError(f'Unknown data augmentation policy "{data_augmentation}"')
+
+        if backend == "pil":
+            # Note: we could just convert to pure tensors even in v2.
+            transforms += [T.ToImage() if use_v2 else T.PILToTensor()]
+
+        transforms += [T.ToDtype(torch.float, scale=True)]
+
+        if use_v2:
+            transforms += [
+                T.ConvertBoundingBoxFormat(tv_tensors.BoundingBoxFormat.XYXY),
+                T.SanitizeBoundingBoxes(),
+                T.ToPureTensor(),
+            ]
+
+        self.transforms = T.Compose(transforms)
+
+    def __call__(self, img, target):
+        return self.transforms(img, target)
+
+
+class DetectionPresetEval:
+    def __init__(self, backend="pil", use_v2=False):
+        T, _ = get_modules(use_v2)
+        transforms = []
+        backend = backend.lower()
+        if backend == "pil":
+            # Note: we could just convert to pure tensors even in v2?
+            transforms += [T.ToImage() if use_v2 else T.PILToTensor()]
+        elif backend == "tensor":
+            transforms += [T.PILToTensor()]
+        elif backend == "tv_tensor":
+            transforms += [T.ToImage()]
+        else:
+            raise ValueError(f"backend can be 'tv_tensor', 'tensor' or 'pil', but got {backend}")
+
+        transforms += [T.ToDtype(torch.float, scale=True)]
+
+        if use_v2:
+            transforms += [T.ToPureTensor()]
+
+        self.transforms = T.Compose(transforms)
+
+    def __call__(self, img, target):
+        return self.transforms(img, target)
diff --git a/tv-detection/retinanet_resnet101_fpn.isc b/tv-detection/retinanet_resnet101_fpn.isc
new file mode 100644
index 00000000..d79bdb8b
--- /dev/null
+++ b/tv-detection/retinanet_resnet101_fpn.isc
@@ -0,0 +1,7 @@
+experiment_name="retinanet_resnet101_fpn"
+gpu_type="24GB VRAM GPU"
+nnodes = 11
+venv_path = "~/.venv/bin/activate"
+output_path = "~/outputs/retinanet_resnet101_fpn"
+command="train_cycling.py --dataset coco --model retinanet_resnet101_fpn --epochs 26 --lr-steps 16 22 -b 2 --aspect-ratio-group-factor 3 --opt adamw --lr 0.001 --data-path=/mnt/.node1/Open-Datasets/coco --resume $OUTPUT_PATH/checkpoint.isc --tboard-path $OUTPUT_PATH/tb"
+
diff --git a/tv-detection/train_cycling.py b/tv-detection/train_cycling.py
new file mode 100644
index 00000000..6861bb84
--- /dev/null
+++ b/tv-detection/train_cycling.py
@@ -0,0 +1,315 @@
+from cycling_utils import TimestampedTimer
+
+timer = TimestampedTimer()
+timer.report('importing Timer')
+
+import os
+from pathlib import Path
+import argparse
+import presets
+import torch
+import torch.utils.data
+import torchvision
+import utils
+from coco_utils import get_coco
+from coco_eval import CocoEvaluator
+from coco_utils import get_coco_api_from_dataset
+
+import torchvision.models.detection
+import torchvision.models.detection.mask_rcnn
+from torchvision.models.detection import MaskRCNN, RetinaNet
+from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
+from engine import evaluate, train_one_epoch
+from group_by_aspect_ratio import create_aspect_ratio_groups
+from torchvision.transforms import InterpolationMode
+from transforms import SimpleCopyPaste
+
+from cycling_utils import InterruptableDistributedSampler, InterruptableDistributedGroupedBatchSampler, MetricsTracker
+
+timer.report('importing everything else')
+
+def copypaste_collate_fn(batch):
+    copypaste = SimpleCopyPaste(blending=True, resize_interpolation=InterpolationMode.BILINEAR)
+    return copypaste(*utils.collate_fn(batch))
+
+def get_dataset(is_train, args):
+    image_set = "train" if is_train else "val"
+    num_classes, mode = {"coco": (91, "instances"), "coco_kp": (2, "person_keypoints")}[args.dataset]
+    with_masks = "mask" in args.model
+    ds = get_coco(
+        root=args.data_path,
+        image_set=image_set,
+        transforms=get_transform(is_train, args),
+        mode=mode,
+        use_v2=args.use_v2,
+        with_masks=with_masks,
+    )
+    return ds, num_classes
+
+def get_transform(is_train, args):
+    if is_train:
+        return presets.DetectionPresetTrain(
+            data_augmentation=args.data_augmentation, backend=args.backend, use_v2=args.use_v2
+        )
+    elif args.weights and args.test_only:
+        weights = torchvision.models.get_weight(args.weights)
+        trans = weights.transforms()
+        return lambda img, target: (trans(img), target)
+    else:
+        return presets.DetectionPresetEval(backend=args.backend, use_v2=args.use_v2)
+    
+def _get_iou_types(model): # intersection over union (iou) types
+    model_without_ddp = model
+    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+        model_without_ddp = model.module
+    iou_types = ["bbox"]
+    if isinstance(model_without_ddp, torchvision.models.detection.MaskRCNN):
+        iou_types.append("segm")
+    if isinstance(model_without_ddp, torchvision.models.detection.KeypointRCNN):
+        iou_types.append("keypoints")
+    return iou_types
+
+timer.report('defined other functions')
+
+def main(args, timer):
+
+    if args.backend.lower() == "tv_tensor" and not args.use_v2:
+        raise ValueError("Use --use-v2 if you want to use the tv_tensor backend.")
+    if args.dataset not in ("coco", "coco_kp"):
+        raise ValueError(f"Dataset should be coco or coco_kp, got {args.dataset}")
+    if "keypoint" in args.model and args.dataset != "coco_kp":
+        raise ValueError("Oops, if you want Keypoint detection, set --dataset coco_kp")
+    if args.dataset == "coco_kp" and args.use_v2:
+        raise ValueError("KeyPoint detection doesn't support V2 transforms yet")
+
+    utils.init_distributed_mode(args)
+    print(args)
+
+    device = torch.device(args.device)
+
+    if args.use_deterministic_algorithms:
+        torch.use_deterministic_algorithms(True)
+
+    timer.report('main preliminaries')
+
+    dataset_train, num_classes = get_dataset(is_train=True, args=args)
+    dataset_test, _ = get_dataset(is_train=False, args=args)
+
+    timer.report('loading data')
+
+    group_ids = create_aspect_ratio_groups(dataset_train, k=args.aspect_ratio_group_factor)
+    train_sampler = InterruptableDistributedGroupedBatchSampler(dataset_train, group_ids, args.batch_size)
+    test_sampler = InterruptableDistributedSampler(dataset_test)
+
+    timer.report('creating data samplers')
+
+    train_collate_fn = utils.collate_fn
+    if args.use_copypaste:
+        if args.data_augmentation != "lsj":
+            raise RuntimeError("SimpleCopyPaste algorithm currently only supports the 'lsj' data augmentation policies")
+        print("Using copypaste_collate_fn for train_collate_fn")
+        train_collate_fn = copypaste_collate_fn
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train, batch_sampler=train_sampler, num_workers=args.workers, collate_fn=train_collate_fn
+    )
+    data_loader_test = torch.utils.data.DataLoader(
+        dataset_test, batch_size=1, sampler=test_sampler, num_workers=args.workers, collate_fn=utils.collate_fn
+    )
+
+    timer.report('creating data loaders')
+
+    kwargs = {"trainable_backbone_layers": args.trainable_backbone_layers}
+    if args.data_augmentation in ["multiscale", "lsj"]:
+        kwargs["_skip_resize"] = True
+    if "rcnn" in args.model:
+        if args.rpn_score_thresh is not None:
+            kwargs["rpn_score_thresh"] = args.rpn_score_thresh
+
+    if args.model == "maskrcnn_resnet101_fpn":
+        backbone  = resnet_fpn_backbone(backbone_name="resnet101", weights="ResNet101_Weights.IMAGENET1K_V1")
+        model = MaskRCNN(backbone=backbone, num_classes=num_classes)
+    elif args.model == "retinanet_resnet50_fpn":
+        backbone  = resnet_fpn_backbone(backbone_name="resnet50", weights="ResNet50_Weights.IMAGENET1K_V1")
+        model = RetinaNet(backbone=backbone, num_classes=num_classes)
+    elif args.model == "retinanet_resnet101_fpn":
+        backbone  = resnet_fpn_backbone(backbone_name="resnet101", weights="ResNet101_Weights.IMAGENET1K_V1")
+        model = RetinaNet(backbone=backbone, num_classes=num_classes)
+    model.to(device)
+
+    timer.report('creating model and .to(device)')
+
+    if args.distributed and args.sync_bn:
+        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+
+    model_without_ddp = model
+    if args.distributed:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
+        model_without_ddp = model.module
+
+    timer.report('preparing model for distributed training')
+
+    if args.norm_weight_decay is None:
+        parameters = [p for p in model.parameters() if p.requires_grad]
+    else:
+        param_groups = torchvision.ops._utils.split_normalization_params(model)
+        wd_groups = [args.norm_weight_decay, args.weight_decay]
+        parameters = [{"params": p, "weight_decay": w} for p, w in zip(param_groups, wd_groups) if p]
+
+    opt_name = args.opt.lower()
+    if opt_name.startswith("sgd"):
+        optimizer = torch.optim.SGD(
+            parameters, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, 
+            nesterov="nesterov" in opt_name,
+        )
+    elif opt_name == "adamw":
+        optimizer = torch.optim.AdamW(parameters, lr=args.lr, weight_decay=args.weight_decay)
+    else:
+        raise RuntimeError(f"Invalid optimizer {args.opt}. Only SGD and AdamW are supported.")
+
+    scaler = torch.cuda.amp.GradScaler() if args.amp else None
+
+    timer.report('optimizer and scaler')
+
+    ## OUTER LR_SCHEDULER
+    args.lr_scheduler = args.lr_scheduler.lower()
+    if args.lr_scheduler == "multisteplr":
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_steps, gamma=args.lr_gamma)
+    elif args.lr_scheduler == "cosineannealinglr":
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
+    else:
+        raise RuntimeError(
+            f"Invalid lr scheduler '{args.lr_scheduler}'. Only MultiStepLR and CosineAnnealingLR are supported."
+        )
+    
+    ## WARMUP LR_SCHEDULER
+    warmup_factor = 1.0 / 1000
+    warmup_iters = min(1000, len(data_loader_train) - 1)
+    warmup_lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+            optimizer, start_factor=warmup_factor, total_iters=warmup_iters
+        )
+
+    timer.report('learning rate schedulers')
+
+    coco = get_coco_api_from_dataset(data_loader_test.dataset)
+    iou_types = _get_iou_types(model)
+    coco_evaluator = CocoEvaluator(coco, iou_types)
+
+    timer.report('init coco evaluator')
+
+    metrics = {"train": MetricsTracker(), "val": MetricsTracker()}
+
+    # RETRIEVE CHECKPOINT
+    Path(args.resume).parent.mkdir(parents=True, exist_ok=True)
+    checkpoint = None
+    if args.resume and os.path.isfile(args.resume): # If we're resuming...
+        print("RESUMING FROM CURRENT JOB")
+        checkpoint = torch.load(args.resume, map_location="cpu")
+    elif args.prev_resume and os.path.isfile(args.prev_resume):
+        print(f"RESUMING FROM PREVIOUS JOB {args.prev_resume}")
+        checkpoint = torch.load(args.prev_resume, map_location="cpu")
+    if checkpoint is not None:
+        args.start_epoch = checkpoint["epoch"]
+        model_without_ddp.load_state_dict(checkpoint["model"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
+        warmup_lr_scheduler.load_state_dict(checkpoint["warmup_lr_scheduler"])
+        train_sampler.load_state_dict(checkpoint["train_sampler"])
+        test_sampler.load_state_dict(checkpoint["test_sampler"])
+        if args.amp:
+            scaler.load_state_dict(checkpoint["scaler"])
+        # Evaluator and metrics
+        coco_evaluator.img_ids = checkpoint["img_ids"]
+        coco_evaluator.eval_imgs = checkpoint["eval_imgs"]
+        metrics = checkpoint["metrics"]
+
+    timer.report('retrieving checkpoint')
+
+    if args.test_only:
+        # We disable the cudnn benchmarking because it can noticeably affect the accuracy
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+        epoch = 0
+        coco_evaluator, timer, metrics = evaluate(
+            model, data_loader_test, epoch, test_sampler, args, coco_evaluator, optimizer, 
+            lr_scheduler, warmup_lr_scheduler, train_sampler, device, scaler, timer, metrics
+        )
+        return
+
+    for epoch in range(args.start_epoch, args.epochs):
+        
+        print('\n')
+        print(f"EPOCH :: {epoch}")
+        print('\n')
+
+        with train_sampler.in_epoch(epoch):
+            timer = TimestampedTimer() # obtain time trial for each epoch
+            model, timer, metrics = train_one_epoch(
+                model, optimizer, data_loader_train, train_sampler, test_sampler, 
+                lr_scheduler, warmup_lr_scheduler, args, device, coco_evaluator, 
+                epoch, scaler, timer, metrics
+            )
+
+            with test_sampler.in_epoch(epoch):
+                timer = TimestampedTimer() # obtain time trial for each epoch
+                coco_evaluator, timer, metrics = evaluate(
+                    model, data_loader_test, epoch, test_sampler, args, coco_evaluator, 
+                    optimizer, lr_scheduler, warmup_lr_scheduler, train_sampler, device, 
+                    scaler, timer, metrics
+                )
+                
+
+def get_args_parser(add_help=True):
+    parser = argparse.ArgumentParser(description="PyTorch Detection Training", add_help=add_help)
+    parser.add_argument("--data-path", default=None, type=str, help="dataset path")
+    parser.add_argument("--dataset",default="coco",type=str, help="dataset name. Use coco for object detection and instance segmentation and coco_kp for Keypoint detection",)
+    parser.add_argument("--model", default="maskrcnn_resnet50_fpn", type=str, help="model name")
+    parser.add_argument("--device", default="cuda", type=str, help="device (Use cuda or cpu Default: cuda)")
+    parser.add_argument("-b", "--batch-size", default=2, type=int, help="images per gpu, the total batch size is $NGPU x batch_size")
+    parser.add_argument("--epochs", default=26, type=int, metavar="N", help="number of total epochs to run")
+    parser.add_argument("-j", "--workers", default=4, type=int, metavar="N", help="number of data loading workers (default: 4)")
+    parser.add_argument("--opt", default="sgd", type=str, help="optimizer")
+    parser.add_argument("--lr", default=0.02, type=float,help="initial learning rate, 0.02 is the default value for training on 8 gpus and 2 images_per_gpu")
+    parser.add_argument("--momentum", default=0.9, type=float, metavar="M", help="momentum")
+    parser.add_argument("--wd","--weight-decay",default=1e-4,type=float,metavar="W",help="weight decay (default: 1e-4)",dest="weight_decay",)
+    parser.add_argument("--norm-weight-decay",default=None,type=float,help="weight decay for Normalization layers (default: None, same value as --wd)")
+    parser.add_argument("--lr-scheduler", default="multisteplr", type=str, help="name of lr scheduler (default: multisteplr)")
+    parser.add_argument("--lr-step-size", default=8, type=int, help="decrease lr every step-size epochs (multisteplr scheduler only)")
+    parser.add_argument("--lr-steps",default=[16, 22],nargs="+",type=int,help="decrease lr every step-size epochs (multisteplr scheduler only)")
+    parser.add_argument("--lr-gamma", default=0.1, type=float, help="decrease lr by a factor of lr-gamma (multisteplr scheduler only)")
+    parser.add_argument("--print-freq", default=1, type=int, help="print frequency")
+    parser.add_argument("--output-dir", default=".", type=str, help="path to save outputs")
+
+    parser.add_argument("--start_epoch", default=0, type=int, help="start epoch")
+    parser.add_argument("--resume", default="", type=str, help="path of checkpoint")
+    parser.add_argument("--prev-resume", default=None, help="path of previous job checkpoint for strong fail resume", dest="prev_resume") # for checkpointing
+    parser.add_argument("--tboard-path", default=None, help="path for saving tensorboard logs", dest="tboard_path") # for checkpointing
+    
+    parser.add_argument("--aspect-ratio-group-factor", default=3, type=int)
+    parser.add_argument("--rpn-score-thresh", default=None, type=float, help="rpn score threshold for faster-rcnn")
+    parser.add_argument("--trainable-backbone-layers", default=None, type=int, help="number of trainable layers of backbone")
+    parser.add_argument("--data-augmentation", default="hflip", type=str, help="data augmentation policy (default: hflip)")
+    parser.add_argument("--sync-bn",dest="sync_bn",help="Use sync batch norm",action="store_true")
+    parser.add_argument("--test-only",dest="test_only",help="Only test the model",action="store_true")
+    parser.add_argument("--use-deterministic-algorithms", action="store_true", help="Forces the use of deterministic algorithms only.")
+
+    # distributed training parameters
+    parser.add_argument("--world-size", default=1, type=int, help="number of distributed processes")
+    parser.add_argument("--dist-url", default="env://", type=str, help="url used to set up distributed training")
+    parser.add_argument("--weights", default=None, type=str, help="the weights enum name to load")
+    parser.add_argument("--weights-backbone", default=None, type=str, help="the backbone weights enum name to load")
+
+    # Mixed precision training parameters
+    parser.add_argument("--amp", action="store_true", help="Use torch.cuda.amp for mixed precision training")
+
+    # Use CopyPaste augmentation training parameter
+    parser.add_argument("--use-copypaste", action="store_true",help="Use CopyPaste data augmentation. Works only with data-augmentation='lsj'.",)
+    parser.add_argument("--backend", default="PIL", type=str.lower, help="PIL or tensor - case insensitive")
+    parser.add_argument("--use-v2", action="store_true", help="Use V2 transforms")
+
+    return parser
+
+
+if __name__ == "__main__":
+    args = get_args_parser().parse_args()
+    main(args, timer)
diff --git a/tv-detection/transforms.py b/tv-detection/transforms.py
new file mode 100644
index 00000000..e07ccfc9
--- /dev/null
+++ b/tv-detection/transforms.py
@@ -0,0 +1,601 @@
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+import torchvision
+from torch import nn, Tensor
+from torchvision import ops
+from torchvision.transforms import functional as F, InterpolationMode, transforms as T
+
+
+def _flip_coco_person_keypoints(kps, width):
+    flip_inds = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
+    flipped_data = kps[:, flip_inds]
+    flipped_data[..., 0] = width - flipped_data[..., 0]
+    # Maintain COCO convention that if visibility == 0, then x, y = 0
+    inds = flipped_data[..., 2] == 0
+    flipped_data[inds] = 0
+    return flipped_data
+
+
+class Compose:
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, image, target):
+        for t in self.transforms:
+            image, target = t(image, target)
+        return image, target
+
+
+class RandomHorizontalFlip(T.RandomHorizontalFlip):
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if torch.rand(1) < self.p:
+            image = F.hflip(image)
+            if target is not None:
+                _, _, width = F.get_dimensions(image)
+                target["boxes"][:, [0, 2]] = width - target["boxes"][:, [2, 0]]
+                if "masks" in target:
+                    target["masks"] = target["masks"].flip(-1)
+                if "keypoints" in target:
+                    keypoints = target["keypoints"]
+                    keypoints = _flip_coco_person_keypoints(keypoints, width)
+                    target["keypoints"] = keypoints
+        return image, target
+
+
+class PILToTensor(nn.Module):
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        image = F.pil_to_tensor(image)
+        return image, target
+
+
+class ToDtype(nn.Module):
+    def __init__(self, dtype: torch.dtype, scale: bool = False) -> None:
+        super().__init__()
+        self.dtype = dtype
+        self.scale = scale
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if not self.scale:
+            return image.to(dtype=self.dtype), target
+        image = F.convert_image_dtype(image, self.dtype)
+        return image, target
+
+
+class RandomIoUCrop(nn.Module):
+    def __init__(
+        self,
+        min_scale: float = 0.3,
+        max_scale: float = 1.0,
+        min_aspect_ratio: float = 0.5,
+        max_aspect_ratio: float = 2.0,
+        sampler_options: Optional[List[float]] = None,
+        trials: int = 40,
+    ):
+        super().__init__()
+        # Configuration similar to https://github.com/weiliu89/caffe/blob/ssd/examples/ssd/ssd_coco.py#L89-L174
+        self.min_scale = min_scale
+        self.max_scale = max_scale
+        self.min_aspect_ratio = min_aspect_ratio
+        self.max_aspect_ratio = max_aspect_ratio
+        if sampler_options is None:
+            sampler_options = [0.0, 0.1, 0.3, 0.5, 0.7, 0.9, 1.0]
+        self.options = sampler_options
+        self.trials = trials
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if target is None:
+            raise ValueError("The targets can't be None for this transform.")
+
+        if isinstance(image, torch.Tensor):
+            if image.ndimension() not in {2, 3}:
+                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")
+            elif image.ndimension() == 2:
+                image = image.unsqueeze(0)
+
+        _, orig_h, orig_w = F.get_dimensions(image)
+
+        while True:
+            # sample an option
+            idx = int(torch.randint(low=0, high=len(self.options), size=(1,)))
+            min_jaccard_overlap = self.options[idx]
+            if min_jaccard_overlap >= 1.0:  # a value larger than 1 encodes the leave as-is option
+                return image, target
+
+            for _ in range(self.trials):
+                # check the aspect ratio limitations
+                r = self.min_scale + (self.max_scale - self.min_scale) * torch.rand(2)
+                new_w = int(orig_w * r[0])
+                new_h = int(orig_h * r[1])
+                aspect_ratio = new_w / new_h
+                if not (self.min_aspect_ratio <= aspect_ratio <= self.max_aspect_ratio):
+                    continue
+
+                # check for 0 area crops
+                r = torch.rand(2)
+                left = int((orig_w - new_w) * r[0])
+                top = int((orig_h - new_h) * r[1])
+                right = left + new_w
+                bottom = top + new_h
+                if left == right or top == bottom:
+                    continue
+
+                # check for any valid boxes with centers within the crop area
+                cx = 0.5 * (target["boxes"][:, 0] + target["boxes"][:, 2])
+                cy = 0.5 * (target["boxes"][:, 1] + target["boxes"][:, 3])
+                is_within_crop_area = (left < cx) & (cx < right) & (top < cy) & (cy < bottom)
+                if not is_within_crop_area.any():
+                    continue
+
+                # check at least 1 box with jaccard limitations
+                boxes = target["boxes"][is_within_crop_area]
+                ious = torchvision.ops.boxes.box_iou(
+                    boxes, torch.tensor([[left, top, right, bottom]], dtype=boxes.dtype, device=boxes.device)
+                )
+                if ious.max() < min_jaccard_overlap:
+                    continue
+
+                # keep only valid boxes and perform cropping
+                target["boxes"] = boxes
+                target["labels"] = target["labels"][is_within_crop_area]
+                target["boxes"][:, 0::2] -= left
+                target["boxes"][:, 1::2] -= top
+                target["boxes"][:, 0::2].clamp_(min=0, max=new_w)
+                target["boxes"][:, 1::2].clamp_(min=0, max=new_h)
+                image = F.crop(image, top, left, new_h, new_w)
+
+                return image, target
+
+
+class RandomZoomOut(nn.Module):
+    def __init__(
+        self, fill: Optional[List[float]] = None, side_range: Tuple[float, float] = (1.0, 4.0), p: float = 0.5
+    ):
+        super().__init__()
+        if fill is None:
+            fill = [0.0, 0.0, 0.0]
+        self.fill = fill
+        self.side_range = side_range
+        if side_range[0] < 1.0 or side_range[0] > side_range[1]:
+            raise ValueError(f"Invalid canvas side range provided {side_range}.")
+        self.p = p
+
+    @torch.jit.unused
+    def _get_fill_value(self, is_pil):
+        # type: (bool) -> int
+        # We fake the type to make it work on JIT
+        return tuple(int(x) for x in self.fill) if is_pil else 0
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if isinstance(image, torch.Tensor):
+            if image.ndimension() not in {2, 3}:
+                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")
+            elif image.ndimension() == 2:
+                image = image.unsqueeze(0)
+
+        if torch.rand(1) >= self.p:
+            return image, target
+
+        _, orig_h, orig_w = F.get_dimensions(image)
+
+        r = self.side_range[0] + torch.rand(1) * (self.side_range[1] - self.side_range[0])
+        canvas_width = int(orig_w * r)
+        canvas_height = int(orig_h * r)
+
+        r = torch.rand(2)
+        left = int((canvas_width - orig_w) * r[0])
+        top = int((canvas_height - orig_h) * r[1])
+        right = canvas_width - (left + orig_w)
+        bottom = canvas_height - (top + orig_h)
+
+        if torch.jit.is_scripting():
+            fill = 0
+        else:
+            fill = self._get_fill_value(F._is_pil_image(image))
+
+        image = F.pad(image, [left, top, right, bottom], fill=fill)
+        if isinstance(image, torch.Tensor):
+            # PyTorch's pad supports only integers on fill. So we need to overwrite the colour
+            v = torch.tensor(self.fill, device=image.device, dtype=image.dtype).view(-1, 1, 1)
+            image[..., :top, :] = image[..., :, :left] = image[..., (top + orig_h) :, :] = image[
+                ..., :, (left + orig_w) :
+            ] = v
+
+        if target is not None:
+            target["boxes"][:, 0::2] += left
+            target["boxes"][:, 1::2] += top
+
+        return image, target
+
+
+class RandomPhotometricDistort(nn.Module):
+    def __init__(
+        self,
+        contrast: Tuple[float, float] = (0.5, 1.5),
+        saturation: Tuple[float, float] = (0.5, 1.5),
+        hue: Tuple[float, float] = (-0.05, 0.05),
+        brightness: Tuple[float, float] = (0.875, 1.125),
+        p: float = 0.5,
+    ):
+        super().__init__()
+        self._brightness = T.ColorJitter(brightness=brightness)
+        self._contrast = T.ColorJitter(contrast=contrast)
+        self._hue = T.ColorJitter(hue=hue)
+        self._saturation = T.ColorJitter(saturation=saturation)
+        self.p = p
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if isinstance(image, torch.Tensor):
+            if image.ndimension() not in {2, 3}:
+                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")
+            elif image.ndimension() == 2:
+                image = image.unsqueeze(0)
+
+        r = torch.rand(7)
+
+        if r[0] < self.p:
+            image = self._brightness(image)
+
+        contrast_before = r[1] < 0.5
+        if contrast_before:
+            if r[2] < self.p:
+                image = self._contrast(image)
+
+        if r[3] < self.p:
+            image = self._saturation(image)
+
+        if r[4] < self.p:
+            image = self._hue(image)
+
+        if not contrast_before:
+            if r[5] < self.p:
+                image = self._contrast(image)
+
+        if r[6] < self.p:
+            channels, _, _ = F.get_dimensions(image)
+            permutation = torch.randperm(channels)
+
+            is_pil = F._is_pil_image(image)
+            if is_pil:
+                image = F.pil_to_tensor(image)
+                image = F.convert_image_dtype(image)
+            image = image[..., permutation, :, :]
+            if is_pil:
+                image = F.to_pil_image(image)
+
+        return image, target
+
+
+class ScaleJitter(nn.Module):
+    """Randomly resizes the image and its bounding boxes  within the specified scale range.
+    The class implements the Scale Jitter augmentation as described in the paper
+    `"Simple Copy-Paste is a Strong Data Augmentation Method for Instance Segmentation" <https://arxiv.org/abs/2012.07177>`_.
+
+    Args:
+        target_size (tuple of ints): The target size for the transform provided in (height, weight) format.
+        scale_range (tuple of ints): scaling factor interval, e.g (a, b), then scale is randomly sampled from the
+            range a <= scale <= b.
+        interpolation (InterpolationMode): Desired interpolation enum defined by
+            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``.
+    """
+
+    def __init__(
+        self,
+        target_size: Tuple[int, int],
+        scale_range: Tuple[float, float] = (0.1, 2.0),
+        interpolation: InterpolationMode = InterpolationMode.BILINEAR,
+        antialias=True,
+    ):
+        super().__init__()
+        self.target_size = target_size
+        self.scale_range = scale_range
+        self.interpolation = interpolation
+        self.antialias = antialias
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if isinstance(image, torch.Tensor):
+            if image.ndimension() not in {2, 3}:
+                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")
+            elif image.ndimension() == 2:
+                image = image.unsqueeze(0)
+
+        _, orig_height, orig_width = F.get_dimensions(image)
+
+        scale = self.scale_range[0] + torch.rand(1) * (self.scale_range[1] - self.scale_range[0])
+        r = min(self.target_size[1] / orig_height, self.target_size[0] / orig_width) * scale
+        new_width = int(orig_width * r)
+        new_height = int(orig_height * r)
+
+        image = F.resize(image, [new_height, new_width], interpolation=self.interpolation, antialias=self.antialias)
+
+        if target is not None:
+            target["boxes"][:, 0::2] *= new_width / orig_width
+            target["boxes"][:, 1::2] *= new_height / orig_height
+            if "masks" in target:
+                target["masks"] = F.resize(
+                    target["masks"],
+                    [new_height, new_width],
+                    interpolation=InterpolationMode.NEAREST,
+                    antialias=self.antialias,
+                )
+
+        return image, target
+
+
+class FixedSizeCrop(nn.Module):
+    def __init__(self, size, fill=0, padding_mode="constant"):
+        super().__init__()
+        size = tuple(T._setup_size(size, error_msg="Please provide only two dimensions (h, w) for size."))
+        self.crop_height = size[0]
+        self.crop_width = size[1]
+        self.fill = fill  # TODO: Fill is currently respected only on PIL. Apply tensor patch.
+        self.padding_mode = padding_mode
+
+    def _pad(self, img, target, padding):
+        # Taken from the functional_tensor.py pad
+        if isinstance(padding, int):
+            pad_left = pad_right = pad_top = pad_bottom = padding
+        elif len(padding) == 1:
+            pad_left = pad_right = pad_top = pad_bottom = padding[0]
+        elif len(padding) == 2:
+            pad_left = pad_right = padding[0]
+            pad_top = pad_bottom = padding[1]
+        else:
+            pad_left = padding[0]
+            pad_top = padding[1]
+            pad_right = padding[2]
+            pad_bottom = padding[3]
+
+        padding = [pad_left, pad_top, pad_right, pad_bottom]
+        img = F.pad(img, padding, self.fill, self.padding_mode)
+        if target is not None:
+            target["boxes"][:, 0::2] += pad_left
+            target["boxes"][:, 1::2] += pad_top
+            if "masks" in target:
+                target["masks"] = F.pad(target["masks"], padding, 0, "constant")
+
+        return img, target
+
+    def _crop(self, img, target, top, left, height, width):
+        img = F.crop(img, top, left, height, width)
+        if target is not None:
+            boxes = target["boxes"]
+            boxes[:, 0::2] -= left
+            boxes[:, 1::2] -= top
+            boxes[:, 0::2].clamp_(min=0, max=width)
+            boxes[:, 1::2].clamp_(min=0, max=height)
+
+            is_valid = (boxes[:, 0] < boxes[:, 2]) & (boxes[:, 1] < boxes[:, 3])
+
+            target["boxes"] = boxes[is_valid]
+            target["labels"] = target["labels"][is_valid]
+            if "masks" in target:
+                target["masks"] = F.crop(target["masks"][is_valid], top, left, height, width)
+
+        return img, target
+
+    def forward(self, img, target=None):
+        _, height, width = F.get_dimensions(img)
+        new_height = min(height, self.crop_height)
+        new_width = min(width, self.crop_width)
+
+        if new_height != height or new_width != width:
+            offset_height = max(height - self.crop_height, 0)
+            offset_width = max(width - self.crop_width, 0)
+
+            r = torch.rand(1)
+            top = int(offset_height * r)
+            left = int(offset_width * r)
+
+            img, target = self._crop(img, target, top, left, new_height, new_width)
+
+        pad_bottom = max(self.crop_height - new_height, 0)
+        pad_right = max(self.crop_width - new_width, 0)
+        if pad_bottom != 0 or pad_right != 0:
+            img, target = self._pad(img, target, [0, 0, pad_right, pad_bottom])
+
+        return img, target
+
+
+class RandomShortestSize(nn.Module):
+    def __init__(
+        self,
+        min_size: Union[List[int], Tuple[int], int],
+        max_size: int,
+        interpolation: InterpolationMode = InterpolationMode.BILINEAR,
+    ):
+        super().__init__()
+        self.min_size = [min_size] if isinstance(min_size, int) else list(min_size)
+        self.max_size = max_size
+        self.interpolation = interpolation
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        _, orig_height, orig_width = F.get_dimensions(image)
+
+        min_size = self.min_size[torch.randint(len(self.min_size), (1,)).item()]
+        r = min(min_size / min(orig_height, orig_width), self.max_size / max(orig_height, orig_width))
+
+        new_width = int(orig_width * r)
+        new_height = int(orig_height * r)
+
+        image = F.resize(image, [new_height, new_width], interpolation=self.interpolation)
+
+        if target is not None:
+            target["boxes"][:, 0::2] *= new_width / orig_width
+            target["boxes"][:, 1::2] *= new_height / orig_height
+            if "masks" in target:
+                target["masks"] = F.resize(
+                    target["masks"], [new_height, new_width], interpolation=InterpolationMode.NEAREST
+                )
+
+        return image, target
+
+
+def _copy_paste(
+    image: torch.Tensor,
+    target: Dict[str, Tensor],
+    paste_image: torch.Tensor,
+    paste_target: Dict[str, Tensor],
+    blending: bool = True,
+    resize_interpolation: F.InterpolationMode = F.InterpolationMode.BILINEAR,
+) -> Tuple[torch.Tensor, Dict[str, Tensor]]:
+
+    # Random paste targets selection:
+    num_masks = len(paste_target["masks"])
+
+    if num_masks < 1:
+        # Such degerante case with num_masks=0 can happen with LSJ
+        # Let's just return (image, target)
+        return image, target
+
+    # We have to please torch script by explicitly specifying dtype as torch.long
+    random_selection = torch.randint(0, num_masks, (num_masks,), device=paste_image.device)
+    random_selection = torch.unique(random_selection).to(torch.long)
+
+    paste_masks = paste_target["masks"][random_selection]
+    paste_boxes = paste_target["boxes"][random_selection]
+    paste_labels = paste_target["labels"][random_selection]
+
+    masks = target["masks"]
+
+    # We resize source and paste data if they have different sizes
+    # This is something we introduced here as originally the algorithm works
+    # on equal-sized data (for example, coming from LSJ data augmentations)
+    size1 = image.shape[-2:]
+    size2 = paste_image.shape[-2:]
+    if size1 != size2:
+        paste_image = F.resize(paste_image, size1, interpolation=resize_interpolation)
+        paste_masks = F.resize(paste_masks, size1, interpolation=F.InterpolationMode.NEAREST)
+        # resize bboxes:
+        ratios = torch.tensor((size1[1] / size2[1], size1[0] / size2[0]), device=paste_boxes.device)
+        paste_boxes = paste_boxes.view(-1, 2, 2).mul(ratios).view(paste_boxes.shape)
+
+    paste_alpha_mask = paste_masks.sum(dim=0) > 0
+
+    if blending:
+        paste_alpha_mask = F.gaussian_blur(
+            paste_alpha_mask.unsqueeze(0),
+            kernel_size=(5, 5),
+            sigma=[
+                2.0,
+            ],
+        )
+
+    # Copy-paste images:
+    image = (image * (~paste_alpha_mask)) + (paste_image * paste_alpha_mask)
+
+    # Copy-paste masks:
+    masks = masks * (~paste_alpha_mask)
+    non_all_zero_masks = masks.sum((-1, -2)) > 0
+    masks = masks[non_all_zero_masks]
+
+    # Do a shallow copy of the target dict
+    out_target = {k: v for k, v in target.items()}
+
+    out_target["masks"] = torch.cat([masks, paste_masks])
+
+    # Copy-paste boxes and labels
+    boxes = ops.masks_to_boxes(masks)
+    out_target["boxes"] = torch.cat([boxes, paste_boxes])
+
+    labels = target["labels"][non_all_zero_masks]
+    out_target["labels"] = torch.cat([labels, paste_labels])
+
+    # Update additional optional keys: area and iscrowd if exist
+    if "area" in target:
+        out_target["area"] = out_target["masks"].sum((-1, -2)).to(torch.float32)
+
+    if "iscrowd" in target and "iscrowd" in paste_target:
+        # target['iscrowd'] size can be differ from mask size (non_all_zero_masks)
+        # For example, if previous transforms geometrically modifies masks/boxes/labels but
+        # does not update "iscrowd"
+        if len(target["iscrowd"]) == len(non_all_zero_masks):
+            iscrowd = target["iscrowd"][non_all_zero_masks]
+            paste_iscrowd = paste_target["iscrowd"][random_selection]
+            out_target["iscrowd"] = torch.cat([iscrowd, paste_iscrowd])
+
+    # Check for degenerated boxes and remove them
+    boxes = out_target["boxes"]
+    degenerate_boxes = boxes[:, 2:] <= boxes[:, :2]
+    if degenerate_boxes.any():
+        valid_targets = ~degenerate_boxes.any(dim=1)
+
+        out_target["boxes"] = boxes[valid_targets]
+        out_target["masks"] = out_target["masks"][valid_targets]
+        out_target["labels"] = out_target["labels"][valid_targets]
+
+        if "area" in out_target:
+            out_target["area"] = out_target["area"][valid_targets]
+        if "iscrowd" in out_target and len(out_target["iscrowd"]) == len(valid_targets):
+            out_target["iscrowd"] = out_target["iscrowd"][valid_targets]
+
+    return image, out_target
+
+
+class SimpleCopyPaste(torch.nn.Module):
+    def __init__(self, blending=True, resize_interpolation=F.InterpolationMode.BILINEAR):
+        super().__init__()
+        self.resize_interpolation = resize_interpolation
+        self.blending = blending
+
+    def forward(
+        self, images: List[torch.Tensor], targets: List[Dict[str, Tensor]]
+    ) -> Tuple[List[torch.Tensor], List[Dict[str, Tensor]]]:
+        torch._assert(
+            isinstance(images, (list, tuple)) and all([isinstance(v, torch.Tensor) for v in images]),
+            "images should be a list of tensors",
+        )
+        torch._assert(
+            isinstance(targets, (list, tuple)) and len(images) == len(targets),
+            "targets should be a list of the same size as images",
+        )
+        for target in targets:
+            # Can not check for instance type dict with inside torch.jit.script
+            # torch._assert(isinstance(target, dict), "targets item should be a dict")
+            for k in ["masks", "boxes", "labels"]:
+                torch._assert(k in target, f"Key {k} should be present in targets")
+                torch._assert(isinstance(target[k], torch.Tensor), f"Value for the key {k} should be a tensor")
+
+        # images = [t1, t2, ..., tN]
+        # Let's define paste_images as shifted list of input images
+        # paste_images = [t2, t3, ..., tN, t1]
+        # FYI: in TF they mix data on the dataset level
+        images_rolled = images[-1:] + images[:-1]
+        targets_rolled = targets[-1:] + targets[:-1]
+
+        output_images: List[torch.Tensor] = []
+        output_targets: List[Dict[str, Tensor]] = []
+
+        for image, target, paste_image, paste_target in zip(images, targets, images_rolled, targets_rolled):
+            output_image, output_data = _copy_paste(
+                image,
+                target,
+                paste_image,
+                paste_target,
+                blending=self.blending,
+                resize_interpolation=self.resize_interpolation,
+            )
+            output_images.append(output_image)
+            output_targets.append(output_data)
+
+        return output_images, output_targets
+
+    def __repr__(self) -> str:
+        s = f"{self.__class__.__name__}(blending={self.blending}, resize_interpolation={self.resize_interpolation})"
+        return s
diff --git a/tv-detection/utils.py b/tv-detection/utils.py
new file mode 100644
index 00000000..dc4298ab
--- /dev/null
+++ b/tv-detection/utils.py
@@ -0,0 +1,281 @@
+import datetime
+import errno
+import os
+import time
+from collections import defaultdict, deque
+
+import torch
+import torch.distributed as dist  
+
+class SmoothedValue:
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median, avg=self.avg, global_avg=self.global_avg, max=self.max, value=self.value
+        )
+
+
+def all_gather(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data: any picklable object
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+    data_list = [None] * world_size
+    dist.all_gather_object(data_list, data)
+    return data_list
+
+
+def reduce_dict(input_dict, average=True):
+    """
+    Args:
+        input_dict (dict): all the values will be reduced
+        average (bool): whether to do average or sum
+    Reduce the values in the dictionary from all processes so that all processes
+    have the averaged results. Returns a dict with the same fields as
+    input_dict, after reduction.
+    """
+    world_size = get_world_size()
+    if world_size < 2:
+        return input_dict
+    with torch.inference_mode():
+        names = []
+        values = []
+        # sort the keys so that they are consistent across processes
+        for k in sorted(input_dict.keys()):
+            names.append(k)
+            values.append(input_dict[k])
+        values = torch.stack(values, dim=0)
+        dist.all_reduce(values)
+        if average:
+            values /= world_size
+        reduced_dict = {k: v for k, v in zip(names, values)}
+    return reduced_dict
+
+
+class MetricLogger:
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{attr}'")
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(f"{name}: {str(meter)}")
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, i, print_freq, header=None):
+        # i = 0
+        if not header:
+            header = ""
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt="{avg:.4f}")
+        data_time = SmoothedValue(fmt="{avg:.4f}")
+        space_fmt = ":" + str(len(str(len(iterable)))) + "d"
+        if torch.cuda.is_available():
+            log_msg = self.delimiter.join(
+                [
+                    header,
+                    "[{0" + space_fmt + "}/{1}]",
+                    "eta: {eta}",
+                    "{meters}",
+                    "time: {time}",
+                    "data: {data}",
+                    "max mem: {memory:.0f}",
+                ]
+            )
+        else:
+            log_msg = self.delimiter.join(
+                [header, "[{0" + space_fmt + "}/{1}]", "eta: {eta}", "{meters}", "time: {time}", "data: {data}"]
+            )
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0: # or i == len(iterable) - 1: ## EDITED - POTENTIALLY UNNECESSARY
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
+                else:
+                    print(
+                        log_msg.format(
+                            i, len(iterable), eta=eta_string, meters=str(self), time=str(iter_time), data=str(data_time)
+                        )
+                    )
+            i += 1 
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print(f"{header} Total time: {total_time_str}") # ({total_time / len(iterable):.4f} s / it)") ## EDITED
+
+
+def collate_fn(batch):
+    return tuple(zip(*batch))
+
+
+def mkdir(path):
+    try:
+        os.makedirs(path)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print("Not using distributed mode")
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = "nccl"
+    print(f"| distributed init (rank {args.rank}): {args.dist_url}", flush=True)
+    torch.distributed.init_process_group(
+        backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank
+    )
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
diff --git a/tv-segmentation/README.md b/tv-segmentation/README.md
index 2c7391c8..7226f371 100644
--- a/tv-segmentation/README.md
+++ b/tv-segmentation/README.md
@@ -4,40 +4,18 @@ This folder contains reference training scripts for semantic segmentation.
 They serve as a log of how to train specific models, as provide baseline
 training and evaluation scripts to quickly bootstrap research.
 
-All models have been trained on 8x V100 GPUs.
+You must ensure all dependencies in "requirements.txt" are installed, and
+run "prep.py" to download pretrained model weights before launching your
+training job.
 
-You must modify the following flags:
+You can run the training routines for the following models using cli.
 
-`--data-path=/path/to/dataset`
-
-`--nproc_per_node=<number_of_gpus_available>`
-
-## fcn_resnet50
-```
-torchrun --nproc_per_node=8 train.py --lr 0.02 --dataset coco -b 4 --model fcn_resnet50 --aux-loss --weights-backbone ResNet50_Weights.IMAGENET1K_V1
-```
-
-## fcn_resnet101
-```
-torchrun --nproc_per_node=8 train.py --lr 0.02 --dataset coco -b 4 --model fcn_resnet101 --aux-loss --weights-backbone ResNet101_Weights.IMAGENET1K_V1
-```
-
-## deeplabv3_resnet50
-```
-torchrun --nproc_per_node=8 train.py --lr 0.02 --dataset coco -b 4 --model deeplabv3_resnet50 --aux-loss --weights-backbone ResNet50_Weights.IMAGENET1K_V1
-```
-
-## deeplabv3_resnet101
-```
-torchrun --nproc_per_node=8 train.py --lr 0.02 --dataset coco -b 4 --model deeplabv3_resnet101 --aux-loss --weights-backbone ResNet101_Weights.IMAGENET1K_V1
-```
-
-## deeplabv3_mobilenet_v3_large
+### RetinaNet
 ```
-torchrun --nproc_per_node=8 train.py --dataset coco -b 4 --model deeplabv3_mobilenet_v3_large --aux-loss --wd 0.000001 --weights-backbone MobileNet_V3_Large_Weights.IMAGENET1K_V1
+isc train ./retinanet_resnet101_fpn.isc
 ```
 
-## lraspp_mobilenet_v3_large
+### Mask R-CNN
 ```
-torchrun --nproc_per_node=8 train.py --dataset coco -b 4 --model lraspp_mobilenet_v3_large --wd 0.000001 --weights-backbone MobileNet_V3_Large_Weights.IMAGENET1K_V1
+isc train ./maskrcnn_resnet101_fpn.isc
 ```
diff --git a/tv-segmentation/deeplabv3_mobilenet_v3_large.isc b/tv-segmentation/deeplabv3_mobilenet_v3_large.isc
index 5dfccfb1..fc5f198f 100644
--- a/tv-segmentation/deeplabv3_mobilenet_v3_large.isc
+++ b/tv-segmentation/deeplabv3_mobilenet_v3_large.isc
@@ -1,7 +1,7 @@
 
 experiment_name="deeplabv3_mobilenet_v3_large"
 gpu_type="24GB VRAM GPU"
-nnodes = 9
+nnodes = 11
 venv_path = "~/.venv/bin/activate"
-output_path = "~/output_tv"
-command="train_cycling.py --dataset coco -b 4 --model deeplabv3_mobilenet_v3_large --aux-loss --wd 0.000001 --weights-backbone MobileNet_V3_Large_Weights.IMAGENET1K_V1 --data-path=/workspace/datasets/coco --resume $OUTPUT_PATH/checkpoint.isc"
+output_path = "~/outputs/deeplabv3_mobilenet_v3_large"
+command="train_cycling.py --dataset coco -b 2 --model deeplabv3_mobilenet_v3_large --lr 0.001 --aux-loss --wd 0.000001 --weights-backbone MobileNet_V3_Large_Weights.IMAGENET1K_V1 --data-path=/mnt/.node1/Open-Datasets/coco --resume $OUTPUT_PATH/checkpoint.isc --tboard-path $OUTPUT_PATH/tb"
diff --git a/tv-segmentation/fcn_resnet101.isc b/tv-segmentation/fcn_resnet101.isc
index 9c91f9cb..b37d26f2 100644
--- a/tv-segmentation/fcn_resnet101.isc
+++ b/tv-segmentation/fcn_resnet101.isc
@@ -1,7 +1,7 @@
 
-experiment_name="seg-fcn_resnet50"
+experiment_name="fcn_resnet101"
 gpu_type="24GB VRAM GPU"
-nnodes = 9
+nnodes = 11
 venv_path = "~/.venv/bin/activate"
-output_path = "~/output_tv"
-command="train_cycling.py --lr 0.0002 --dataset coco -b 8 --model fcn_resnet101 --aux-loss --weights-backbone ResNet101_Weights.IMAGENET1K_V1 --data-path=/workspace/datasets/coco --epochs=1 --resume $OUTPUT_PATH/checkpoint.isc"
+output_path = "~/outputs/fcn_resnet101"
+command="train_cycling.py --lr 0.001 --dataset coco -b 4 --model fcn_resnet101 --aux-loss --weights-backbone ResNet101_Weights.IMAGENET1K_V1 --data-path=/mnt/.node1/Open-Datasets/coco --resume $OUTPUT_PATH/checkpoint.isc --tboard-path $OUTPUT_PATH/tb --prev-resume /mnt/Client/StrongHumans/strong_adam/outputs/fcn_resnet101/13a04288-f836-44eb-aa0c-3e5dda18ddb5/checkpoint.isc"
\ No newline at end of file
diff --git a/tv-segmentation/prep.py b/tv-segmentation/prep.py
index de0974a8..6a39b7fc 100644
--- a/tv-segmentation/prep.py
+++ b/tv-segmentation/prep.py
@@ -1,5 +1,9 @@
 import torchvision
+from torchvision.models import resnet101, ResNet101_Weights, mobilenet_v3_large, MobileNet_V3_Large_Weights
 
 weights = torchvision.models.get_weight('MobileNet_V3_Large_Weights.IMAGENET1K_V1')
 weights = torchvision.models.get_weight('ResNet101_Weights.IMAGENET1K_V1')
 
+_ = resnet101(weights=ResNet101_Weights.IMAGENET1K_V1)
+_ = mobilenet_v3_large(weights=MobileNet_V3_Large_Weights.IMAGENET1K_V1)
+
diff --git a/tv-segmentation/train_cycling.py b/tv-segmentation/train_cycling.py
index 013a88d7..d1ebbc7c 100644
--- a/tv-segmentation/train_cycling.py
+++ b/tv-segmentation/train_cycling.py
@@ -1,9 +1,11 @@
-import datetime
-import os
-import time
-import warnings
+from cycling_utils import TimestampedTimer
+
+timer = TimestampedTimer()
+timer.report('importing Timer')
 
+import os
 from pathlib import Path
+import argparse
 import presets
 import torch
 import torch.utils.data
@@ -13,101 +15,72 @@
 from torch import nn
 from torch.optim.lr_scheduler import PolynomialLR
 from torchvision.transforms import functional as F, InterpolationMode
-from cycling_utils import InterruptableDistributedSampler, atomic_torch_save
+from cycling_utils import InterruptableDistributedSampler, atomic_torch_save, MetricsTracker
+
+from torch.utils.tensorboard import SummaryWriter
 
+timer.report('importing everything else')
 
 def get_dataset(dir_path, name, image_set, transform):
     def sbd(*args, **kwargs):
         return torchvision.datasets.SBDataset(*args, mode="segmentation", **kwargs)
-
     paths = {
         "voc": (dir_path, torchvision.datasets.VOCSegmentation, 21),
         "voc_aug": (dir_path, sbd, 21),
         "coco": (dir_path, get_coco, 21),
     }
     p, ds_fn, num_classes = paths[name]
-
     ds = ds_fn(p, image_set=image_set, transforms=transform)
     return ds, num_classes
 
-
 def get_transform(train, args):
     if train:
         return presets.SegmentationPresetTrain(base_size=520, crop_size=480)
     elif args.weights and args.test_only:
         weights = torchvision.models.get_weight(args.weights)
         trans = weights.transforms()
-
         def preprocessing(img, target):
             img = trans(img)
             size = F.get_dimensions(img)[1:]
             target = F.resize(target, size, interpolation=InterpolationMode.NEAREST)
             return img, F.pil_to_tensor(target)
-
         return preprocessing
     else:
         return presets.SegmentationPresetEval(base_size=520)
 
-
 def criterion(inputs, target):
     losses = {}
     for name, x in inputs.items():
         losses[name] = nn.functional.cross_entropy(x, target, ignore_index=255)
-
     if len(losses) == 1:
         return losses["out"]
-
     return losses["out"] + 0.5 * losses["aux"]
 
+def train_one_epoch(
+        args, model, criterion, optimizer, data_loader_train,
+        train_sampler, test_sampler, confmat, lr_scheduler, 
+        device, epoch, scaler=None, timer=None, metrics=None
+    ):
 
-def evaluate(model, data_loader, device, num_classes):
-    model.eval()
-    confmat = utils.ConfusionMatrix(num_classes)
-    metric_logger = utils.MetricLogger(delimiter="  ")
-    header = "Test:"
-    num_processed_samples = 0
-    with torch.inference_mode():
-        for image, target in metric_logger.log_every(data_loader, 100, header):
-            image, target = image.to(device), target.to(device)
-            output = model(image)
-            output = output["out"]
-
-            confmat.update(target.flatten(), output.argmax(1).flatten())
-            # FIXME need to take into account that the datasets
-            # could have been padded in distributed setup
-            num_processed_samples += image.shape[0]
+    model.train()
 
-        confmat.reduce_from_all_processes()
+    train_step = train_sampler.progress // data_loader_train.batch_size
+    total_steps = len(train_sampler) // data_loader_train.batch_size
+    print(f'\nTraining / resuming epoch {epoch} from training step {train_step}\n')
+    timer.report('launch training routine')
 
-    num_processed_samples = utils.reduce_across_processes(num_processed_samples)
-    if (
-        hasattr(data_loader.dataset, "__len__")
-        and len(data_loader.dataset) != num_processed_samples
-        and torch.distributed.get_rank() == 0
-    ):
-        # See FIXME above
-        warnings.warn(
-            f"It looks like the dataset has {len(data_loader.dataset)} samples, but {num_processed_samples} "
-            "samples were used for the validation, which might bias the results. "
-            "Try adjusting the batch size and / or the world size. "
-            "Setting the world size to 1 is always a safe bet."
-        )
+    for images, target in data_loader_train:
 
-    return confmat
+        images, target = images.to(device), target.to(device)
+        timer.report(f'Epoch: {epoch} batch {train_step}: moving batch data to device')
 
+        optimizer.zero_grad()
 
-def train_one_epoch(model, criterion, optimizer, data_loader, sampler: InterruptableDistributedSampler, lr_scheduler, device, epoch, print_freq, scaler=None):
-    model.train()
-    metric_logger = utils.MetricLogger(delimiter="  ")
-    metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value}"))
-    header = f"Epoch: [{epoch}]"
-    for image, target in metric_logger.log_every(data_loader, sampler.progress // data_loader.batch_size, print_freq, header):
-        image, target = image.to(device), target.to(device)
         with torch.cuda.amp.autocast(enabled=scaler is not None):
-            output = model(image)
+            output = model(images)
             loss = criterion(output, target)
+        timer.report(f'Epoch: {epoch} batch {train_step}: forward pass')
 
-        optimizer.zero_grad()
         if scaler is not None:
             scaler.scale(loss).backward()
             scaler.step(optimizer)
@@ -115,32 +88,121 @@ def train_one_epoch(model, criterion, optimizer, data_loader, sampler: Interrupt
         else:
             loss.backward()
             optimizer.step()
-
         lr_scheduler.step()
 
-        metric_logger.update(loss=loss.item(), lr=optimizer.param_groups[0]["lr"])
+        timer.report(f'Epoch: {epoch} batch {train_step}: backward pass')
 
-        sampler.advance(len(image))
+        metrics["train"].update({"images_seen": len(images), "loss": loss.item()})
+        metrics["train"].reduce() # Reduce to sync metrics between nodes for this batch
+        batch_loss = metrics["train"].local["loss"] / metrics["train"].local["images_seen"]
+        print(f"EPOCH: [{epoch}], BATCH: [{train_step}/{total_steps}], loss: {batch_loss}")
+        metrics["train"].reset_local()
 
-        step = sampler.progress // data_loader.batch_size
-        if utils.is_main_process() and step % 5 == 0:
-            print(f"Saving checkpoint at step {step}")
+        print(f"Saving checkpoint at epoch {epoch} train batch {train_step}")
+        train_sampler.advance(len(images))
+
+        train_step = train_sampler.progress // data_loader_train.batch_size
+        if train_step == total_steps:
+            metrics["train"].end_epoch()
+
+        if utils.is_main_process() and train_step % 1 == 0: # Checkpointing every batch
+
+            writer = SummaryWriter(log_dir=args.tboard_path)
+            writer.add_scalar("Train/loss", batch_loss, train_step + epoch * total_steps)
+            writer.flush()
+            writer.close()
+            
             checkpoint = {
+                "args": args,
+                "epoch": epoch,
                 "model": model.module.state_dict(),
                 "optimizer": optimizer.state_dict(),
                 "lr_scheduler": lr_scheduler.state_dict(),
-                "epoch": epoch,
-                "args": args,
-                "sampler": sampler.state_dict(),
+                "train_sampler": train_sampler.state_dict(),
+                "test_sampler": test_sampler.state_dict(),
+                "confmat": confmat.mat,
+                "confmat_temp": confmat.temp_mat,
+                "metrics": metrics,
             }
+
             if args.amp:
                 checkpoint["scaler"] = scaler.state_dict()
-            atomic_torch_save(checkpoint, args.resume)
+            timer = atomic_torch_save(checkpoint, args.resume, timer)
+
+    return model, timer, metrics
+
+def evaluate(
+        args, model, data_loader_test, num_classes, confmat,
+        optimizer, lr_scheduler, train_sampler, test_sampler, 
+        device, epoch, scaler=None, timer=None, metrics=None,
+    ):
+
+    model.eval()
+
+    test_step = test_sampler.progress // data_loader_test.batch_size
+    print(f'\nEvaluating / resuming epoch {epoch} from eval step {test_step}\n')
+    timer.report('launch evaluation routine')
+
+    with torch.inference_mode():
+
+        for images, target in data_loader_test:
+
+            images, target = images.to(device), target.to(device)
+            timer.report(f'Epoch {epoch} batch: {test_step} moving to device')
+
+            output = model(images)
+            output = output["out"]
+            timer.report(f'Epoch {epoch} batch: {test_step} forward through model')
+
+            confmat.update(target.flatten().detach().cpu(), output.argmax(1).flatten().detach().cpu())
+            confmat.reduce_from_all_processes()
+
+            timer.report(f'Epoch {epoch} batch: {test_step} confmat update')
+
+            print(f"Saving checkpoint at epoch {epoch} eval batch {test_step}")
+            test_sampler.advance(len(images))
+            test_step = test_sampler.progress // data_loader_test.batch_size
+
+            if utils.is_main_process() and test_step % 1 == 0: # Checkpointing every batch
+                checkpoint = {
+                    "args": args,
+                    "epoch": epoch,
+                    "model": model.module.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "lr_scheduler": lr_scheduler.state_dict(),
+                    "train_sampler": train_sampler.state_dict(),
+                    "test_sampler": test_sampler.state_dict(),
+                    "confmat": confmat.mat, # For storing eval metric
+                    "confmat_temp": confmat.temp_mat, # For storing eval metric
+                    "metrics": metrics,
+                }
+                if args.amp:
+                    checkpoint["scaler"] = scaler.state_dict()
+                timer = atomic_torch_save(checkpoint, args.resume, timer)
+
+    # Report key performance metrics
+    acc_global, acc, iu = confmat.compute()
+    acc = acc_global.item() * 100
+    mean_iou = iu.mean().item() * 100
+    metrics["val"].update({"acc": acc, "mean_iou": mean_iou})
+    metrics["val"].reduce()
+    metrics["val"].reset_local()
+    print(f"EPOCH: [{epoch}] EVAL :: acc: {acc:.2f}, mean_iou: {mean_iou:.2f}")
+    confmat.reset()
+
+    if utils.is_main_process():
+        writer = SummaryWriter(log_dir=args.tboard_path)
+        writer.add_scalar("Val/acc", acc, epoch)
+        writer.add_scalar("Val/mean_iou", mean_iou, epoch)
+        writer.flush()
+        writer.close()
 
+    return confmat, timer, metrics
 
-def main(args):
-    if args.output_dir:
-        utils.mkdir(args.output_dir)
+timer.report('defined other functions')
+
+
+def main(args, timer):
 
     utils.init_distributed_mode(args)
     print(args)
@@ -155,37 +217,37 @@ def main(args):
     else:
         torch.backends.cudnn.benchmark = True
 
-    dataset, num_classes = get_dataset(args.data_path, args.dataset, "train", get_transform(True, args))
+    timer.report('main preliminaries')
+
+    dataset_train, num_classes = get_dataset(args.data_path, args.dataset, "train", get_transform(True, args))
     dataset_test, _ = get_dataset(args.data_path, args.dataset, "val", get_transform(False, args))
 
-    # if args.distributed:
-    train_sampler = InterruptableDistributedSampler(dataset)
-    test_sampler = torch.utils.data.distributed.DistributedSampler(dataset_test, shuffle=False)
-    # else:
-    #     train_sampler = torch.utils.data.RandomSampler(dataset)
-    #     test_sampler = torch.utils.data.SequentialSampler(dataset_test)
-
-    data_loader = torch.utils.data.DataLoader(
-        dataset,
-        batch_size=args.batch_size,
-        sampler=train_sampler,
-        num_workers=args.workers,
-        collate_fn=utils.collate_fn,
-        drop_last=True,
-    )
+    timer.report('loading data')
+
+    train_sampler = InterruptableDistributedSampler(dataset_train)
+    test_sampler = InterruptableDistributedSampler(dataset_test)
+
+    timer.report('creating data samplers')
 
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train, batch_size=args.batch_size, sampler=train_sampler, num_workers=args.workers, 
+        collate_fn=utils.collate_fn, drop_last=True,
+    )
     data_loader_test = torch.utils.data.DataLoader(
-        dataset_test, batch_size=1, sampler=test_sampler, num_workers=args.workers, collate_fn=utils.collate_fn
+        dataset_test, batch_size=1, sampler=test_sampler, num_workers=args.workers, 
+        collate_fn=utils.collate_fn
     )
 
+    timer.report('creating data loaders')
+
     model = torchvision.models.get_model(
-        args.model,
-        weights=args.weights,
-        weights_backbone=args.weights_backbone,
-        num_classes=num_classes,
+        args.model, weights=args.weights, weights_backbone=args.weights_backbone, num_classes=num_classes, 
         aux_loss=args.aux_loss,
     )
     model.to(device)
+
+    timer.report('creating model and .to(device)')
+
     if args.distributed:
         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
 
@@ -194,10 +256,13 @@ def main(args):
         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
         model_without_ddp = model.module
 
+    timer.report('preparing model for distributed training')
+
     params_to_optimize = [
         {"params": [p for p in model_without_ddp.backbone.parameters() if p.requires_grad]},
         {"params": [p for p in model_without_ddp.classifier.parameters() if p.requires_grad]},
     ]
+
     if args.aux_loss:
         params = [p for p in model_without_ddp.aux_classifier.parameters() if p.requires_grad]
         params_to_optimize.append({"params": params, "lr": args.lr * 10})
@@ -205,7 +270,9 @@ def main(args):
 
     scaler = torch.cuda.amp.GradScaler() if args.amp else None
 
-    iters_per_epoch = len(data_loader)
+    timer.report('optimizer and scaler')
+
+    iters_per_epoch = len(data_loader_train)
     main_lr_scheduler = PolynomialLR(
         optimizer, total_iters=iters_per_epoch * (args.epochs - args.lr_warmup_epochs), power=0.9
     )
@@ -231,69 +298,91 @@ def main(args):
     else:
         lr_scheduler = main_lr_scheduler
 
+    timer.report('learning rate schedulers')
+
+    # Init global confmat for eval - eval accumulator
+    confmat = utils.ConfusionMatrix(num_classes)
+    # Init general purpose metrics tracker
+    metrics = {"train": MetricsTracker(), "val": MetricsTracker()}
+
+    timer.report('init metrics')
+
+    # RETRIEVE CHECKPOINT
     Path(args.resume).parent.mkdir(parents=True, exist_ok=True)
-    if args.resume and os.path.isfile(args.resume):
+    checkpoint = None
+    if args.resume and os.path.isfile(args.resume): # If we're resuming...
         checkpoint = torch.load(args.resume, map_location="cpu")
+        print("RESUMING FROM CURRENT JOB")
+    elif args.prev_resume and os.path.isfile(args.prev_resume):
+        print(f"RESUMING FROM PREVIOUS JOB {args.prev_resume}")
+        checkpoint = torch.load(args.prev_resume, map_location="cpu")
+    if checkpoint is not None:
+
+        args.start_epoch = checkpoint["epoch"]
         model_without_ddp.load_state_dict(checkpoint["model"], strict=not args.test_only)
-        if not args.test_only:
-            optimizer.load_state_dict(checkpoint["optimizer"])
-            lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
-            args.start_epoch = checkpoint["epoch"] #+ 1
-            if args.amp:
-                scaler.load_state_dict(checkpoint["scaler"])
-            train_sampler.load_state_dict(checkpoint["sampler"])
+        
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
+        train_sampler.load_state_dict(checkpoint["train_sampler"])
+        test_sampler.load_state_dict(checkpoint["test_sampler"])
+        if args.amp: # Could align this syntactically...
+            scaler.load_state_dict(checkpoint["scaler"])
+        confmat.mat = checkpoint["confmat"]
+        confmat.temp_mat = checkpoint["confmat_temp"]
+        metrics = checkpoint["metrics"]
+            
+    timer.report('retrieving checkpoint')
 
     if args.test_only:
         # We disable the cudnn benchmarking because it can noticeably affect the accuracy
         torch.backends.cudnn.benchmark = False
         torch.backends.cudnn.deterministic = True
-        confmat = evaluate(model, data_loader_test, device=device, num_classes=num_classes)
+        epoch = 0
+        confmat, timer, metrics = evaluate(
+            args, model, data_loader_test, num_classes, confmat,
+            optimizer, lr_scheduler, train_sampler, test_sampler, 
+            device, epoch, scaler, timer, metrics,
+        )
         print(confmat)
         return
 
-    start_time = time.time()
     for epoch in range(args.start_epoch, args.epochs):
-        # if args.distributed:
+
+        print('\n')
+        print(f"EPOCH :: {epoch}")
+        print('\n')
+
         with train_sampler.in_epoch(epoch):
-            train_one_epoch(model, criterion, optimizer, data_loader, train_sampler, lr_scheduler, device, epoch, args.print_freq, scaler)
-            confmat = evaluate(model, data_loader_test, device=device, num_classes=num_classes)
-            print(confmat)
-            if utils.is_main_process():
-                checkpoint = {
-                    "model": model.module.state_dict(),
-                    "optimizer": optimizer.state_dict(),
-                    "lr_scheduler": lr_scheduler.state_dict(),
-                    "epoch": epoch,
-                    "args": args,
-                    "sampler": train_sampler.state_dict(),
-                }
-                if args.amp:
-                    checkpoint["scaler"] = scaler.state_dict()
-                atomic_torch_save(checkpoint, args.resume)
 
-    total_time = time.time() - start_time
-    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
-    print(f"Training time {total_time_str}")
+            timer = TimestampedTimer() # obtain time trial for each epoch
+            model, timer, metrics = train_one_epoch(
+                args, model, criterion, optimizer, data_loader_train,
+                train_sampler, test_sampler, confmat, lr_scheduler, 
+                device, epoch, scaler, timer, metrics
+            )
+            timer.report(f'training for epoch {epoch}')
 
+            with test_sampler.in_epoch(epoch):
+                
+                timer = TimestampedTimer() # obtain time trial for each epoch
+                confmat, timer, metrics = evaluate(
+                    args, model, data_loader_test, num_classes, confmat,
+                    optimizer, lr_scheduler, train_sampler, test_sampler, 
+                    device, epoch, scaler, timer, metrics,
+                )
+                timer.report(f'evaluation for epoch {epoch}')
 
-def get_args_parser(add_help=True):
-    import argparse
 
+def get_args_parser(add_help=True):
     parser = argparse.ArgumentParser(description="PyTorch Segmentation Training", add_help=add_help)
-
     parser.add_argument("--data-path", default="/datasets01/COCO/022719/", type=str, help="dataset path")
     parser.add_argument("--dataset", default="coco", type=str, help="dataset name")
     parser.add_argument("--model", default="fcn_resnet101", type=str, help="model name")
     parser.add_argument("--aux-loss", action="store_true", help="auxiliary loss")
     parser.add_argument("--device", default="cuda", type=str, help="device (Use cuda or cpu Default: cuda)")
-    parser.add_argument(
-        "-b", "--batch-size", default=8, type=int, help="images per gpu, the total batch size is $NGPU x batch_size"
-    )
+    parser.add_argument("-b", "--batch-size", default=8, type=int, help="images per gpu, the total batch size is $NGPU x batch_size", dest="batch_size")
     parser.add_argument("--epochs", default=30, type=int, metavar="N", help="number of total epochs to run")
-
-    parser.add_argument(
-        "-j", "--workers", default=16, type=int, metavar="N", help="number of data loading workers (default: 16)"
-    )
+    parser.add_argument("-j", "--workers", default=16, type=int, metavar="N", help="number of data loading workers (default: 16)")
     parser.add_argument("--lr", default=0.01, type=float, help="initial learning rate")
     parser.add_argument("--momentum", default=0.9, type=float, metavar="M", help="momentum")
     parser.add_argument(
@@ -308,9 +397,11 @@ def get_args_parser(add_help=True):
     parser.add_argument("--lr-warmup-epochs", default=0, type=int, help="the number of epochs to warmup (default: 0)")
     parser.add_argument("--lr-warmup-method", default="linear", type=str, help="the warmup method (default: linear)")
     parser.add_argument("--lr-warmup-decay", default=0.01, type=float, help="the decay for lr")
-    parser.add_argument("--print-freq", default=10, type=int, help="print frequency")
+    parser.add_argument("--print-freq", default=1, type=int, help="print frequency")
     parser.add_argument("--output-dir", default=".", type=str, help="path to save outputs")
     parser.add_argument("--resume", type=str, help="path of checkpoint", required=True)
+    parser.add_argument("--prev-resume", default=None, help="path of previous job checkpoint for strong fail resume", dest="prev_resume") # for checkpointing
+    parser.add_argument("--tboard-path", default=None, help="path for saving tensorboard logs", dest="tboard_path") # for checkpointing
     parser.add_argument("--start-epoch", default=0, type=int, metavar="N", help="start epoch")
     parser.add_argument(
         "--test-only",
@@ -322,7 +413,7 @@ def get_args_parser(add_help=True):
         "--use-deterministic-algorithms", action="store_true", help="Forces the use of deterministic algorithms only."
     )
     # distributed training parameters
-    parser.add_argument("--world-size", default=1, type=int, help="number of distributed processes")
+    parser.add_argument("--world-size", default=9, type=int, help="number of distributed processes")
     parser.add_argument("--dist-url", default="env://", type=str, help="url used to set up distributed training")
 
     parser.add_argument("--weights", default=None, type=str, help="the weights enum name to load")
@@ -336,4 +427,4 @@ def get_args_parser(add_help=True):
 
 if __name__ == "__main__":
     args = get_args_parser().parse_args()
-    main(args)
+    main(args, timer)
diff --git a/tv-segmentation/utils.py b/tv-segmentation/utils.py
index 73bb882d..8f524d1e 100644
--- a/tv-segmentation/utils.py
+++ b/tv-segmentation/utils.py
@@ -7,7 +7,6 @@
 import torch
 import torch.distributed as dist
 
-
 class SmoothedValue:
     """Track a series of values and provide access to smoothed values over a
     window or the global series average.
@@ -66,18 +65,22 @@ def __str__(self):
 class ConfusionMatrix:
     def __init__(self, num_classes):
         self.num_classes = num_classes
-        self.mat = None
+        # temp_mat will accumulate results from images seen on this node
+        self.temp_mat = torch.zeros((num_classes, num_classes), dtype=torch.int64, device='cpu', requires_grad=False)
+        # mat will then store the accumulation of all temp_mats, avoiding multiple-counting
+        self.mat = torch.zeros((num_classes, num_classes), dtype=torch.int64, device='cpu', requires_grad=False)
 
     def update(self, a, b):
         n = self.num_classes
-        if self.mat is None:
-            self.mat = torch.zeros((n, n), dtype=torch.int64, device=a.device)
+        # if self.mat is None:
+        #     self.mat = torch.zeros((n, n), dtype=torch.int64, device=a.device)
         with torch.inference_mode():
             k = (a >= 0) & (a < n)
             inds = n * a[k].to(torch.int64) + b[k]
-            self.mat += torch.bincount(inds, minlength=n**2).reshape(n, n)
+            self.temp_mat += torch.bincount(inds, minlength=n**2).reshape(n, n)
 
     def reset(self):
+        self.temp_mat.zero_()
         self.mat.zero_()
 
     def compute(self):
@@ -88,7 +91,11 @@ def compute(self):
         return acc_global, acc, iu
 
     def reduce_from_all_processes(self):
-        reduce_across_processes(self.mat)
+        reduce_across_processes(self.temp_mat)
+        # add the accumulated results from all nodes
+        self.mat += self.temp_mat.clone()
+        # reset temp_mat
+        self.temp_mat.zero_()
 
     def __str__(self):
         acc_global, acc, iu = self.compute()
@@ -111,7 +118,7 @@ def update(self, **kwargs):
                 v = v.item()
             if not isinstance(v, (float, int)):
                 raise TypeError(
-                    f"This method expects the value of the input arguments to be of type float or int, instead  got {type(v)}"
+                    f"Method expects value of the input arguments to be of type float or int, instead  got {type(v)}"
                 )
             self.meters[k].update(v)
 
@@ -296,5 +303,5 @@ def reduce_across_processes(val):
 
     t = torch.tensor(val, device="cuda")
     dist.barrier()
-    dist.all_reduce(t)
+    dist.all_reduce(t) # default all_reduce op is SUM
     return t