deuglify mnist example

PiperOrigin-RevId: 688922829
google-research · Oct 23, 2024 · c815431 · c815431
1 parent ebb9144
commit c815431
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/examples/mnist_autoencoder.py b/examples/mnist_autoencoder.py
@@ -87,8 +87,7 @@ def get_config():
 
 
 def _make_ds(training: bool):
-  Tfds = kd.data.py.Tfds
-  return Tfds(
+  return kd.data.py.Tfds(
       name="mnist",
       split="train" if training else "test",
       shuffle=True if training else False,

diff --git a/kauldron/data/py/base.py b/kauldron/data/py/base.py
@@ -170,7 +170,7 @@ def ds_for_current_process(self, rng: random.PRNGKey) -> grain.MapDataset:
 
 def _get_num_workers(num_workers: int) -> int:
   """Set the number of workers."""
-  if epy.is_notebook():  # in colab worker_count has to be 0
+  if epy.is_notebook() or epy.is_test():  # in colab worker_count has to be 0
     # TODO(klausg): autodetect if Kernel supports multiprocessing
     # Could check
     # from multiprocessing import spawn