import tensorflow as tf

						pred_label = tf.argmax(distillation_loss["st_logits"], axis=-1, output_type=tf.int32)
						correct = tf.equal(
							tf.cast(tf.ones_like(label_ids, dtype=tf.int32), tf.int32),
							tf.cast(pred_label, tf.int32)
