MWE copy tensor bug

copy_tensor_bug.py

import numpy as np
import os
import tensorflow as tf
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.optimizers import Adam

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'


def get_node(idx):
    idx = idx + 1

    return idx * np.ones((100, 100, 3)), 1


def get_node_pyfunc(idx):
    tensors = tf.py_function(
        func=get_node,
        inp=[idx],
        Tout=(tf.int64, tf.int64)
    )

    tensors[0].set_shape(tf.TensorShape([None, None, None]))
    tensors[1].set_shape(tf.TensorShape([]))

    return tensors


dataset = tf.data.Dataset.range(1000)
dataset = dataset.repeat(20)
dataset = dataset.map(get_node_pyfunc, num_parallel_calls=4)
dataset = dataset.batch(batch_size=20)

model = tf.keras.Sequential()
model.add(Flatten(input_shape=(100, 100, 3)))
model.add(Dense(1000))
model.add(Dense(1000))
model.add(Dense(1))
model.compile(optimizer=Adam(), loss='binary_crossentropy')
model.fit(dataset, verbose=True, steps_per_epoch=1000 / 20, epochs=20)