Tell autograph to not convert data loaders, do not force gym versions to improve downstream imports.

PiperOrigin-RevId: 316699386

Author: Lukasz Kaiser

setup.py

@@ -50,7 +50,7 @@
         'gin-config',
         'google-api-python-client',
         'gunicorn',
-        'gym==0.14.0',
+        'gym',
         'h5py',
         'kfac',
         'mesh-tensorflow',

tensor2tensor/data_generators/problem.py

@@ -399,6 +399,7 @@ def set_task_id(self, new_task_id):
   # END SUBCLASS INTERFACE
   # ============================================================================
 
+  @tf.autograph.experimental.do_not_convert()
   def preprocess(self, dataset, mode, hparams, interleave=True):
     """Runtime preprocessing on the whole dataset.
 
@@ -587,6 +588,7 @@ def maybe_reverse_and_copy(self, example):
     self.maybe_copy_features(example)
     return example
 
+  @tf.autograph.experimental.do_not_convert()
   def dataset(self,
               mode,
               data_dir=None,