` TypeError: int() argument must be a string, a bytes-like object or a number ` raised when fitting a multi input Keras model

Question

I'm currently building a U-net model handling multiple input streams of data with Keras/Tensorflow's Functional API. Even though my model compiles, it raises a TypeError when I try to fit it. This seems to be an issue with the way I handle my inputs. I tried to go back to the simplest model I could think of (using multiple inputs), consulted documentation and tried to find similar models online to compare my implementation, but nothing has worked so far. I am able to make it work with one input stream, however even the simplest multi-input model still raises the following:

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

With the following traceback:

I'm currently passing noise to the model for simplicity (in the shape of the actual data). I tried to pass input shapes instead and call the Input() method within my Unet class instead but this did not solve the issue. Here is the simplified version of the code I'm currently working with:

Import statements

import numpy as np
import datetime
import random
import tensorflow as tf
from tensorflow.keras.callbacks import ModelCheckpoint
from tensorflow.keras.models import *
from tensorflow.keras.layers import *
from tensorflow.keras.layers import Dense, Dropout, Activation, Flatten, Conv2D, MaxPooling2D
from tensorflow import keras
from keras import backend
from tensorflow.keras.optimizers import *

Custom loss function (doesn't seem to be the issue)

class custom_loss:
    def wrapper(self):

        def masked_loss(Y_true, Y_pred):
            '''
                This loss function takes the prediction Y_pred, applies the mask to the first two channels
                which correspond to the real and imag part of the visibilities, then computes the 'chi^2'
                of those two channels.
            '''

            '''
            Uncomment if 0 -> flag and 1 -> not flagged
            
            ones = tf.ones_like(Y_true[:,:,:,2])
            #invert the mask, we want only the masked areas to enter the chi^2
            mask_array = ones - Y_true[:,:,:,2]
            '''
            #mask_array = Y_true[:,:,:,2]
            #mask_array = tf.ones_like(Y_true[:,:,:,2]) - Y_true[:,:,:,2]
            mask_array = Y_true[:,:,:,2]

            Y_pred_real = tf.math.multiply(tf.cast(Y_pred[:,:,:,0] , tf.float64) , tf.cast(mask_array[:,:,:], tf.float64))
            Y_true_real= tf.math.multiply(tf.cast(Y_true[:,:,:,0], tf.float64) , tf.cast(mask_array[:,:,:], tf.float64))
            
            Y_pred_imag = tf.math.multiply(tf.cast(Y_pred[:,:,:,1], tf.float64), tf.cast(mask_array[:,:,:], tf.float64))
            Y_true_imag = tf.math.multiply(tf.cast(Y_true[:,:,:,1], tf.float64), tf.cast(mask_array[:,:,:], tf.float64))
            
            
            ground_truth_reconstructed = tf.complex(Y_true_real, Y_true_imag)
            
            predictions_reconstructed = tf.complex(Y_pred_real, Y_pred_imag)
            
            chi = ground_truth_reconstructed - predictions_reconstructed
            
            chi2 = tf.math.conj(chi) * chi
            
            return tf.math.real(tf.math.reduce_sum(chi2))

        return masked_loss

Simplified model

class Unet: 
    '''
        Multi-input U-net architecture where Conv2D layers are replaced by custom DSS ones
 for inpainting of flagged radio astronomy measurements
    '''
    def __init__(self, inputs, loss_class, weights_dir):
        # inherit from Functional API
        super(Unet, self).__init__()
        self.input_a, self.input_b = inputs
        self.weights_dir = weights_dir
        
        masked_loss  = loss_class.wrapper()

        output_a = self.input_a
        output_b = self.input_b
        
        self.model = Model(inputs, [output_a, output_b])
        self.model._name = "U-NET-V0.6"
        
        #----------- handling missing weights -----------#
        if self.weights_dir != None:
            try:
                print('Loading weights' )
                self.model.load_weights(self.weights_dir)
            except:
                pass
                print('No weights to load' , flush = True)
        
        #--------------- compiling the model -------------#
        self.model.compile(optimizer = Adam(learning_rate = 1e-4), loss = masked_loss , metrics = [ masked_loss ])

        self.model.summary()
        #keras.utils.plot_model(CNN, "multi_input_and_output_model.png")

Getting data and creating/compiling the model

# Generating sample data to test the network
x_test = Input(np.random.normal(0, 1, size = (512, 512, 3)).shape, name="data-val")

x_train_a = Input(np.random.normal(0, 1, size = (512, 512, 3)).shape, name="data_a")
x_train_b = Input(np.random.normal(0, 1, size = (512, 512, 3)).shape, name="data_b")

y_train_a = Input(np.random.normal(0, 1, size = (512, 512, 3)).shape, name="labels_a")
y_train_b = Input(np.random.normal(0, 1, size = (512, 512, 3)).shape, name="labels_b")

# Creating path where networl progress is saved
checkpoint_path = '../latest.hdf5'

# Creating an instance of the loss class
loss = custom_loss()

modelcheckpoint = ModelCheckpoint(save_best_only = True, save_weights_only = True, verbose = 1, filepath = checkpoint_path, monitor = 'val_loss')

# Creating a model with the CNN
CNN_obj = Unet([x_train_a, x_train_b], loss, checkpoint_path)

Fitting the model (i.e. where the error is raised)

# Fitting the model
CNN_obj.model.fit([x_train_a, x_train_b], [y_train, y_train], batch_size = 5, epochs = 2)

Any pointers, adequate documentation or potential solutions to this issue would be greatly appreciated.

Thank you for your time! (: