Colab notebook enformer-training.ipynb does not iterate properly

When I run the provided notebook "enformer-training.ipynb" on colab I run into the following error when I include the basenji_barnyard-dataset as stated here #367:

OutOfRangeError Traceback (most recent call last)
[/usr/local/lib/python3.10/dist-packages/tensorflow/python/data/ops/iterator_ops.py](https://localhost:8080/#) in __next__(self)
796 try:
--> 797 return self._next_internal()
798 except errors.OutOfRangeError:

4 frames
[/usr/local/lib/python3.10/dist-packages/tensorflow/python/data/ops/iterator_ops.py](https://localhost:8080/#) in _next_internal(self)
779 with context.execution_mode(context.SYNC):
--> 780 ret = gen_dataset_ops.iterator_get_next(
781 self._iterator_resource,

[/usr/local/lib/python3.10/dist-packages/tensorflow/python/ops/gen_dataset_ops.py](https://localhost:8080/#) in iterator_get_next(iterator, output_types, output_shapes, name)
3015 except _core._NotOkStatusException as e:
-> 3016 _ops.raise_from_not_ok_status(e, name)
3017 except _core._FallbackException:

[/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py](https://localhost:8080/#) in raise_from_not_ok_status(e, name)
7261 e.message += (" name: " + name if name is not None else "")
-> 7262 raise core._status_to_exception(e) from None # pylint: disable=protected-access
7263

OutOfRangeError: {{function_node __wrapped__IteratorGetNext_output_types_2_device_/job:localhost/replica:0/task:0/device:CPU:0}} End of sequence [Op:IteratorGetNext]

During handling of the above exception, another exception occurred:

StopIteration Traceback (most recent call last)
[](https://localhost:8080/#) in ()
1 it = iter(mouse_dataset)
----> 2 example = next(it)

[/usr/local/lib/python3.10/dist-packages/tensorflow/python/data/ops/iterator_ops.py](https://localhost:8080/#) in __next__(self)
797 return self._next_internal()
798 except errors.OutOfRangeError:
--> 799 raise StopIteration
800
801 @property

StopIteration:

Is this a tensorflow version issue or is the dataset I am using not the correct one? Please help!