import tensorflow as tf

            self.assertEqual(logits_.shape, (batch_size,))

    def test_regression(self):
        """Test the type of regression output."""
        batch_size = 8

        hparams = {
            "pretrained_model_name": None,
            "regr_strategy": "cls_time"
        }
        inputs = tf.placeholder(tf.int32, shape=[batch_size, 6])
        regressor = XLNetRegressor(hparams=hparams)
        logits = regressor(inputs)

        with self.test_session() as sess:
            sess.run(tf.global_variables_initializer())
            logits_ = sess.run(
                logits,
                feed_dict={inputs: np.random.randint(30521,
                                                     size=(batch_size, 6))})
            self.assertEqual(logits_.dtype, np.float32)


if __name__ == "__main__":
    tf.test.main()
