diff --git a/deeptrack/datasets/endothelial_vs/endothelial_vs.py b/deeptrack/datasets/endothelial_vs/endothelial_vs.py
index 2a5bb13e0..eb7818fac 100644
--- a/deeptrack/datasets/endothelial_vs/endothelial_vs.py
+++ b/deeptrack/datasets/endothelial_vs/endothelial_vs.py
@@ -2,6 +2,7 @@
 
 import tensorflow_datasets as tfds
 import numpy as np
+import tensorflow as tf
 
 _DESCRIPTION = """
 """
@@ -34,10 +35,10 @@ def _info(self) -> tfds.core.DatasetInfo:
             features=tfds.features.FeaturesDict(
                 {
                     "image": tfds.features.Image(
-                        shape=(None, None, 1), dtype="uint16"
+                        shape=(None, None, 1), dtype=tf.uint16
                     ),
                     "label": tfds.features.Image(
-                        shape=(None, None, 1), dtype="uint16"
+                        shape=(None, None, 1), dtype=tf.uint16
                     ),
                 }
             ),
diff --git a/deeptrack/features.py b/deeptrack/features.py
index ad4998cbb..14beec481 100644
--- a/deeptrack/features.py
+++ b/deeptrack/features.py
@@ -1307,8 +1307,8 @@ def _process_properties(self, propertydict) -> dict:
 
         return propertydict
 
-    def get(self, image, key, **kwargs):
-        return self.collection[key](image)
+    def get(self, image, key, _ID=(), **kwargs):
+        return self.collection[key](image, _ID=_ID)
 
 
 class OneOfDict(Feature):
@@ -1338,8 +1338,8 @@ def _process_properties(self, propertydict) -> dict:
 
         return propertydict
 
-    def get(self, image, key, **kwargs):
-        return self.collection[key](image)
+    def get(self, image, key, _ID=(), **kwargs):
+        return self.collection[key](image, _ID=_ID)
 
 
 # class Dataset(Feature):
@@ -2109,7 +2109,7 @@ def _check_volumes_non_overlapping(self, volume_1, volume_2, min_distance):
         )
 
     def _resample_volume_position(self, volume):
-        """ Draws a new position for the volume. """
+        """Draws a new position for the volume."""
 
         for pdict in volume.properties:
             if "position" in pdict and "_position_sampler" in pdict:
@@ -2122,4 +2122,4 @@ def _resample_volume_position(self, volume):
 
 
 # Alias
-Dataset = TensorflowDataset
\ No newline at end of file
+Dataset = TensorflowDataset