Internal changes

PiperOrigin-RevId: 305956008

official/vision/image_classification/classifier_trainer_test.py

@@ -148,7 +148,7 @@ class ClassifierTest(tf.test.TestCase, parameterized.TestCase):
     tf.io.gfile.rmtree(self.get_temp_dir())
 
   @combinations.generate(distribution_strategy_combinations())
-  def test_end_to_end_train_and_eval_export(self, distribution, model, dataset):
+  def test_end_to_end_train_and_eval(self, distribution, model, dataset):
     """Test train_and_eval and export for Keras classifier models."""
     # Some parameters are not defined as flags (e.g. cannot run
     # classifier_train.py --batch_size=...) by design, so use
@@ -164,38 +164,25 @@ class ClassifierTest(tf.test.TestCase, parameterized.TestCase):
         '--mode=train_and_eval',
     ]
 
-    export_params = basic_params_override()
-    export_path = os.path.join(model_dir, 'export')
-    export_params['export'] = {}
-    export_params['export']['destination'] = export_path
-    export_flags = base_flags + [
-        '--mode=export_only',
-        get_params_override(export_params)
-    ]
-
     run = functools.partial(classifier_trainer.run,
                             strategy_override=distribution)
     run_end_to_end(main=run,
                    extra_flags=train_and_eval_flags,
                    model_dir=model_dir)
-    run_end_to_end(main=run,
-                   extra_flags=export_flags,
-                   model_dir=model_dir)
-    self.assertTrue(os.path.exists(export_path))
 
   @combinations.generate(
       combinations.combine(
-      distribution=[
-          strategy_combinations.one_device_strategy_gpu,
-      ],
-      model=[
-          'efficientnet',
-          'resnet',
-      ],
-      mode='eager',
-      dataset='imagenet',
-      dtype='float16',
-  ))
+          distribution=[
+              strategy_combinations.one_device_strategy_gpu,
+          ],
+          model=[
+              'efficientnet',
+              'resnet',
+          ],
+          mode='eager',
+          dataset='imagenet',
+          dtype='float16',
+      ))
   def test_gpu_train(self, distribution, model, dataset, dtype):
     """Test train_and_eval and export for Keras classifier models."""
     # Some parameters are not defined as flags (e.g. cannot run