Пользовательский слой нормализации в Tensorflow

Question

Я пытаюсь реализовать пользовательский слой нормализации для Tensorflow (2.0). Слой должен рассчитать построчный максимум и затем вернуть нормированные входные данные: x_norm = x / max(x).

Поэтому я следовал документации по созданию пользовательских слоев в TF и ​​нашел решение:

import numpy as np
import tensorflow as tf
from sklearn.datasets import make_regression
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense, Input


class NormalizationLayer(tf.keras.layers.Layer):

    def __init__(self):
        super(NormalizationLayer, self).__init__()

    def build(self, input_shape):
        batch_size = input_shape[0]
        self.max = self.add_weight(name='Max',
                                   shape=(batch_size, ),
                                   trainable=False)

    def call(self, inputs, *args, **kwargs):
        self.max.assign_add(tf.math.reduce_max(inputs, axis=1))
        return tf.transpose(tf.math.divide_no_nan(tf.transpose(inputs), self.max))

Используя следующий простой пример, кажется, что слой делает именно то, что должен:

x, y = make_regression(n_samples=10000, n_features=10, bias=0, noise=1)
x += 10

# Benchmark
x_max = np.max(x, axis=1)
x_norm = (x.T / x_max).T

# Layer
x_tf = tf.constant(x)
norm_layer = NormalizationLayer()
x_tf_norm = norm_layer(x_tf)
x_tf_norm = x_tf_norm.numpy()

print(x_norm[0, 0:5])
print(x_tf_norm[0, 0:5])

Но при использовании того же слоя в последовательном режиме он завершается с ошибкой TypeError:

model = Sequential()
model.add(Input(shape=(10, )))
model.add(NormalizationLayer())
model.add(Dense(16))
model.add(Dense(1))

model.compile(optimizer='adam', loss='mae', metrics=['mae'])

TypeError: int() argument must be a string, a bytes-like object or a number, not 'NoneType'

Любые комментарии приветствуются.