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

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):

'''
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
'''

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))



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

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:
except:
pass
print('No weights to load' , flush = True)

#--------------- compiling the model -------------#

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! (: