internal changes

PiperOrigin-RevId: 721750228
google-research · Jan 31, 2025 · ac06058 · ac06058
1 parent 8806d8d
commit ac06058
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/scenic/dataset_lib/big_transfer/bit.py b/scenic/dataset_lib/big_transfer/bit.py
@@ -82,10 +82,14 @@ def get_dataset(*,
   target_is_onehot = 'onehot' in dataset_configs.pp_train
 
   def pp_fn(x, how):
-    pp = builder.get_preprocess_fn(how)
+    pp = builder.get_preprocess_fn(
+        how, remove_tpu_dtypes=dataset_configs.get('remove_tpu_dtypes', True))
     example = pp(x)
     # to scenic format
-    return {'inputs': example['image'], 'label': example['labels']}
+    return_dict = {'inputs': example['image'], 'label': example['labels']}
+    if dataset_configs.dataset == 'imagenet2012' and 'file_name' in example:
+      return_dict['file_name'] = example['file_name']
+    return return_dict
 
   # E.g. for testing with TAP.
   shuffle_buffer_size = (1000 if num_shards == 1 else

diff --git a/scenic/dataset_lib/big_transfer/builder.py b/scenic/dataset_lib/big_transfer/builder.py
@@ -23,7 +23,7 @@
 
 TPU_SUPPORTED_DTYPES = [
     tf.bool, tf.int32, tf.int64, tf.bfloat16, tf.float32, tf.complex64,
-    tf.uint32
+    tf.uint32,
 ]