diff --git a/darwin/future/data_objects/properties.py b/darwin/future/data_objects/properties.py
index 27635cf6b..f60b917f3 100644
--- a/darwin/future/data_objects/properties.py
+++ b/darwin/future/data_objects/properties.py
@@ -84,8 +84,6 @@ class FullProperty(DefaultDarwin):
     options: Optional[List[PropertyValue]] = None
     granularity: PropertyGranularity = PropertyGranularity("section")
 
-    # model_config = ConfigDict(use_enum_values=True)
-
     def to_create_endpoint(
         self,
     ) -> dict:
diff --git a/darwin/torch/transforms.py b/darwin/torch/transforms.py
index fd60cb3c5..e783d4947 100644
--- a/darwin/torch/transforms.py
+++ b/darwin/torch/transforms.py
@@ -368,8 +368,11 @@ def _pre_process(self, image: np.ndarray, annotation: dict) -> dict:
         if (
             masks is not None and masks.numel() > 0
         ):  # using numel() to check if tensor is non-empty
-            print("WE GOT MASKS")
-            albumentation_dict["masks"] = masks.numpy()
+            if isinstance(masks, torch.Tensor):
+                masks = masks.numpy()
+            if masks.ndim == 3:  # Ensure masks is a list of numpy arrays
+                masks = [masks[i] for i in range(masks.shape[0])]
+            albumentation_dict["masks"] = masks
 
         return albumentation_dict