Hi everyone,
I’m trying to implement a simple feed-forward neural network with a modification based on this paper arxiv:2007.11207 2 (See figure 3b).
Basically the inputs (x) are scaled by a vector, e.g.,  (1,2,3) producing new inputs (1x, 2x, 3x). These are then passed into 3 sub-networks that are independent of each other. The outputs of these networks are then summed together at the end to produce an output.
I have the following TF2 code below but when I try to plot the model using tf.keras.utils.plot_model(model) I get the following error.
AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute '_keras_history'

The code I am using to build the model is:
(for simplicity the sub-networks just have a single layer)
import tensorflow as tf

def build_model():
    input_shape = 1
    output_shape = 1
    units = 128
    scales = tf.convert_to_tensor([[1],[2],[3]], dtype=tf.float32)

    input_layer = tf.keras.Input(shape=(input_shape,))

    # create sub-networks
    xs = []
    for scale in scales:
        scaled_input = tf.keras.layers.multiply([input_layer, scale])
        xs.append(tf.keras.layers.Dense(units, activation='relu')(scaled_input))

    # add the outputs of each sub-network
    output_layer = tf.keras.layers.add([x for x in xs])
    output_layer = tf.keras.layers.Dense(output_shape, activation="linear")(output_layer)

    model = tf.keras.models.Model(inputs=input_layer, outputs=output_layer)

    return model

model = build_model()

# then the following throws the error.
tf.keras.utils.plot_model(model)

Thanks in advance for any help!