diff --git a/uncertainty_baselines/datasets/toxic_comments.py b/uncertainty_baselines/datasets/toxic_comments.py
index a89be1fe..0c7af199 100644
--- a/uncertainty_baselines/datasets/toxic_comments.py
+++ b/uncertainty_baselines/datasets/toxic_comments.py
@@ -252,11 +252,14 @@ def _download_and_prepare(self, dl_manager, download_config=None):
     return self._tfds_dataset_builder._download_and_prepare(  # pylint: disable=protected-access
         dl_manager, download_config)
 
-  def _as_dataset(self,
-                  split: tfds.Split,
-                  decoders=None,
-                  read_config=None,
-                  shuffle_files=False) -> tf.data.Dataset:
+  def _as_dataset(
+      self,
+      split: tfds.Split,
+      decoders=None,
+      read_config=None,
+      shuffle_files=False,
+      file_format: str | tfds.core.FileFormat | None = None,
+  ) -> tf.data.Dataset:
     raise NotImplementedError
 
   # Note that we override `as_dataset` instead of `_as_dataset` to avoid any