Update xception_test with correct target modality

PiperOrigin-RevId: 187658067

Author: Ryan Sepassi

tensor2tensor/models/xception_test.py

@@ -32,7 +32,7 @@
 
 class XceptionTest(tf.test.TestCase):
 
-  def _testXception(self, img_size, output_size):
+  def _testXception(self, img_size):
     vocab_size = 9
     batch_size = 3
     x = np.random.random_integers(
@@ -42,6 +42,7 @@ def _testXception(self, img_size, output_size):
     hparams = xception.xception_tiny()
     p_hparams = problem_hparams.test_problem_hparams(vocab_size, vocab_size)
     p_hparams.input_modality["inputs"] = (registry.Modalities.IMAGE, None)
+    p_hparams.target_modality = (registry.Modalities.CLASS_LABEL, vocab_size)
     with self.test_session() as session:
       features = {
           "inputs": tf.constant(x, dtype=tf.int32),
@@ -51,13 +52,13 @@ def _testXception(self, img_size, output_size):
       logits, _ = model(features)
       session.run(tf.global_variables_initializer())
       res = session.run(logits)
-    self.assertEqual(res.shape, output_size + (1, vocab_size))
+    self.assertEqual(res.shape, (batch_size, 1, 1, 1, vocab_size))
 
-  def testXceptionSmall(self):
-    self._testXception(img_size=9, output_size=(3, 5, 5))
+  def testXceptionSmallImage(self):
+    self._testXception(img_size=9)
 
-  def testXceptionLarge(self):
-    self._testXception(img_size=256, output_size=(3, 8, 8))
+  def testXceptionLargeImage(self):
+    self._testXception(img_size=256)
 
 
 if __name__ == "__main__":