import tensorflow as tf

      out = s(buffers, transitions, training=True)
      self.assertEqual(tf.float32, out.dtype)
      self.assertEqual((1, embedding_dims), out.shape)

  def testSNLIClassifierAndTrainer(self):
    with tf.device(self._test_device):
      vocab_size = 40
      batch_size = 2
      d_embed = 10
      sequence_length = 15
