Tensorflow 2.0 stack () вызывает ошибку неинициализированных тензоров

Question

Я пишу пользовательский слой, где мне нужно l oop через размер пакета, а затем размер изображения rgb. Я все еще пытаюсь понять, как Tensorflow реализует циклы for, и я не уверен, что это связано с ошибкой, которую я здесь представляю.

Вот некоторый псевдокод:

    @tf.function()
    def _crop_and_resize(self, imgs, boxes, to_size):
        # prepare kernel_h and kernel_w

        n_images = tf.shape(imgs)[0]
        outputs = tf.TensorArray(dtype=tf.float32, size=n_images)
        for i in tf.range(n_images):
            # in the call to _bilinear we enter the inner loop
            output = self._bilinear(
                kernel_h[i],
                kernel_w[i],
                imgs[i])
            outputs.write(i, output)
        return outputs.stack()


    def _bilinear(self, kernel_h, kernel_w, img):
        n_channels = tf.shape(img)[2]
        result_channels = tf.TensorArray(dtype=tf.float32, size=n_channels)
        for i in tf.range(n_channels):
            result_channels.write(i,
                tf.matmul(
                    tf.matmul(kernel_h, tf.tile(img[:, :, i], [1, 1])),
                    kernel_w, transpose_b=True))
        return tf.transpose(result_channels.stack(), perm=[1,2,0])

Я получаю следующую ошибку:

InvalidArgumentError: Пытался составить список, который содержит только неинициализированные тензоры и имеет не полностью определенный element_shape: [?,?,?] [[{{ node model_17 / att_1 / PartitionedCall / TensorArrayV2Stack / TensorListStack}}]] [Op: __ inference_distributed_function_11150] Стек вызова функции: распределенная_функция

Я видел много примеров использования TensorArray и stack в этом для одиночного for-l oop, но я не уверен, что мое гнездо для l oop вызывает проблему.