--- a
+++ b/libs/datasets/collate_batch.py
@@ -0,0 +1,26 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+from utils.image_list import to_image_list
+
+class BatchCollator(object):
+    """
+    From a list of samples from the dataset,
+    returns the batched images and targets.
+    This should be passed to the DataLoader
+    """
+
+    def __init__(self, size_divisible=0, df_used=False, boundary=False):
+        self.size_divisible = size_divisible
+        self.df_used = df_used
+        self.boundary = boundary
+
+    def __call__(self, batch):
+        transposed_batch = list(zip(*batch))
+        images = to_image_list(transposed_batch[0], self.size_divisible)
+        targets = to_image_list(transposed_batch[1], self.size_divisible)
+        
+        dfs = to_image_list(transposed_batch[2], self.size_divisible) if self.df_used else None
+        
+        dist_maps = to_image_list(transposed_batch[3], self.size_divisible) if self.boundary else None
+
+        return images, targets, dfs, dist_maps
+