Создание подклассов модели Tensorflow 2.0 Keras

Я пытаюсь реализовать простую UNet-подобную модель, используя метод подкласса модели. Вот мой код:

import tensorflow as tf 
from tensorflow import keras as K

class Enc_block(K.layers.Layer):
    def __init__(self, in_dim):
        super(Enc_block, self).__init__()
        self.conv_layer =  K.layers.SeparableConv2D(in_dim,3, padding='same', activation='relu')
        self.batchnorm_layer = K.layers.BatchNormalization()
        self.pool_layer = K.layers.SeparableConv2D(in_dim,3, padding='same',strides=2, activation='relu')

    def call(self, x):
        x = self.conv_layer(x)
        x = self.batchnorm_layer(x)
        x = self.conv_layer(x)
        x = self.batchnorm_layer(x)
        return self.pool_layer(x), x


class Dec_block(K.layers.Layer):
    def __init__(self, in_dim):
        super(Dec_block, self).__init__()
        self.conv_layer =  K.layers.SeparableConv2D(in_dim,3, padding='same', activation='relu')
        self.batchnorm_layer = K.layers.BatchNormalization()

    def call(self, x):
        x = self.conv_layer(x)
        x = self.batchnorm_layer(x)
        x = self.conv_layer(x)
        x = self.batchnorm_layer(x)
        return x

class Bottleneck(K.layers.Layer):
    def __init__(self, in_dim):
        super(Bottleneck, self).__init__()
        self.conv_1layer =  K.layers.SeparableConv2D(in_dim,1, padding='same', activation='relu')
        self.conv_3layer =  K.layers.SeparableConv2D(in_dim,3, padding='same', activation='relu')
        self.batchnorm_layer = K.layers.BatchNormalization()

    def call(self, x):
        x = self.conv_1layer(x)
        x = self.batchnorm_layer(x)
        x = self.conv_3layer(x)
        x = self.batchnorm_layer(x)
        return x

class Output_block(K.layers.Layer):
    def __init__(self, in_dim):
        super(Output_block, self).__init__()
        self.logits = K.layers.SeparableConv2D(in_dim,3, padding='same', activation=None)
        self.out = K.layers.Softmax()

    def call(self, x):
        x_logits = self.logits(x)
        x = self.out(x_logits)
        return x_logits, x

class UNetModel(K.Model):
    def __init__(self,in_dim):
        super(UNetModel, self).__init__()
        self.encoder_block = Enc_block(in_dim)
        self.bottleneck = Bottleneck(in_dim)
        self.decoder_block = Dec_block(in_dim)
        self.output_block = Output_block(in_dim)


    def call(self, inputs, training=None):
        x, x_skip1 = self.encoder_block(32)(inputs)
        x, x_skip2 = self.encoder_block(64)(x)
        x, x_skip3 = self.encoder_block(128)(x)
        x, x_skip4 = self.encoder_block(256)(x)
        x = self.bottleneck(x)
        x = K.layers.UpSampling2D(size=(2,2))(x)
        x = K.layers.concatenate([x,x_skip4],axis=-1)
        x = self.decoder_block(256)(x)
        x = K.layers.UpSampling2D(size=(2,2))(x) #56x56
        x = K.layers.concatenate([x,x_skip3],axis=-1)
        x = self.decoder_block(128)(x)
        x = K.layers.UpSampling2D(size=(2,2))(x) #112x112
        x = K.layers.concatenate([x,x_skip2],axis=-1)
        x = self.decoder_block(64)(x)
        x = K.layers.UpSampling2D(size=(2,2))(x) #224x224
        x = K.layers.concatenate([x,x_skip1],axis=-1)
        x = self.decoder_block(32)(x)
        x_logits, x = self.output_block(2)(x)
        return x_logits, x

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

ValueError: Input 0 of layer separable_conv2d is incompatible with the layer: expected ndim=4, found ndim=0. Full shape received: []

Я не уверен, что это правильный способ реализации сети в tf.keras. Идея заключалась в том, чтобы реализовать блоки кодировщика и декодера путем создания подклассов слоев keras и подкласса модели позже.