Как решить проблему с batch_size?

def construct_autoregressive_mask(target):
    """
        Args: Original target of word ids, shape [batch, seq_len]
        Returns: a mask of shape [batch, seq_len, seq_len].
    """
    batch_size, seq_len = target.shape.as_list()
    
    print(batch_size)

    tri_matrix = np.zeros((seq_len, seq_len))
    tri_matrix[np.tril_indices(seq_len)] = 1
    
    mask = tf.convert_to_tensor(tri_matrix, dtype=tf.float32)

    masks = tf.tile(tf.expand_dims(mask, 0), (batch_size, 1, 1))  # copies
    return masks