def METHOD_NAME(self, cost_type, with_length):
    strategy = tf.distribute.MirroredStrategy(devices=["/cpu:0"])
    attention_probs = tf.random.uniform([2, 5, 6])
    gold_alignment = tf.random.uniform([2, 5, 6])
    if with_length:
        sequence_length = tf.constant([4, 5], dtype=tf.int32)
    else:
        sequence_length = None
    with strategy.scope():
        losses.guided_alignment_cost(
            attention_probs,
            gold_alignment,
            sequence_length=sequence_length,
            cost_type=cost_type,
        )