How to train this model using tensorflow functional API?

[aliencaocao]

xInput = Input(input_shape, dtype=tf.uint8)
x = Resizing(224, 224)(xInput)
x = norm_layer(xInput)
x = tf.cast(x, tf.float32)
from perceiver import Perceiver
perceiver = Perceiver(
    input_channels = 3,          # number of channels for each token of the input
    input_axis = 2,              # number of axis for input data (2 for images, 3 for video)
    num_freq_bands = 6,          # number of freq bands, with original value (2 * K + 1)
    max_freq = 10.,              # maximum frequency, hyperparameter depending on how fine the data is
    depth = 6,                   # depth of net
    num_latents = 256,           # number of latents
    latent_dim = 512,            # latent dimension
    cross_heads = 1,             # number of heads for cross attention. paper said 1
    latent_heads = 8,            # number of heads for latent self attention, 8
    cross_dim_head = 64,
    latent_dim_head = 64,
    num_classes = 13,          # output number of classes
    attn_dropout = 0.,
    ff_dropout = 0.,
)
xOutput = perceiver(x)

I tried to ran this, but it keeps throwing this error:

TypeError: in user code:

    c:\program files\python38\lib\site-packages\perceiver\perceiver.py:101 call  *
        enc_pos = repeat(enc_pos, "... -> b ...", b=b)
    c:\program files\python38\lib\site-packages\einops\einops.py:461 repeat  *
        return reduce(tensor, pattern, reduction='repeat', **axes_lengths)
    c:\program files\python38\lib\site-packages\einops\einops.py:368 reduce  *
        return recipe.apply(tensor)
    c:\program files\python38\lib\site-packages\einops\einops.py:204 apply  *
        init_shapes, reduced_axes, axes_reordering, added_axes, final_shapes = self.reconstruct_from_shape(
    c:\program files\python38\lib\site-packages\einops\einops.py:192 reconstruct_from_shape  **
        final_shapes.append(_product(lengths))
    c:\program files\python38\lib\site-packages\einops\einops.py:22 _product
        result *= element

    TypeError: unsupported operand type(s) for *=: 'int' and 'NoneType'

Calling model.fit also give the same error above.
The prediction demo you showed works