`_UserObject` object has no attribute `call_and return_conditional_losses` - tensorflow2.0

Question

I am trying to use multistep training. The prediction of the first neural network is used as input for the second neural network. So I need to load the first neural network and call prediction while using the second neural network. I need the first neural network's graph /function and use it in the second neural network.
#neural network class
class Linear(Model):
def __init__(self, n_layers, activate = "tanh", dtype = "float32"):
super(Linear, self).__init__()
self.activation = activate
self.title=title
self.model = Sequential()
self.n_layers=n_layers
#self.model._set_inputs(inputs)
num_layers = len(n_layers)
for l in range(0, num_layers-1):
name = "layer_" + str(l)
m = Dense(n_layers[l+1], input_shape=(n_layers[l],), activation = self.activation, name = name, dtype=dtype)
self.model.add(m)
name = "layer_" + str(num_layers-1)
m = Dense(n_layers[-1], input_shape=(n_layers[-2],), name = name, dtype=dtype)
self.model.add(m)
#tf.function
def __call__(self, X):
Y = self.model(X)
#Pred=self.model.predict(X)
return Y
#Saving and loading
def save(self,name):
tf.saved_model.save(self.model_NN,'saved_model/')
def load(self,name):
restored_saved_model=keras.models.load_model('saved_model/')
return restored_saved_model
#I am training models one after another and using pretection of first model as input of second. So I need to load and save.
def train(self, epoch,multistepping):
for l in range(0, self.number_of_NN):
print("making neural network object",l)
self.NN_list.append(self.net)
for l in range(0,self.number_of_NN):
model = self.pde(self.NN_list[l],self.D,self.dt,self.q, self.cond_i,self.cond_b, self.lr, self.lr_schedule, self.dtype)
print(" training model number ",l)
model.train(epoch[l])
model.save(l)
model.title=l
#model_old=copy.copy(model)
#Bring it outside
def function(xy):
#load weights from previous model
if l>0:
temp_model=model.load(l-1)
x_1 = xy[:, 0][:, None]
y_1 = xy[:, 1][:, None]
U = temp_model.predict(x_1, y_1)
U=U[:,-1]
return np.asarray(U)
if multistepping==1:
self.cond_i.u_func = function
#The error message I am prompted
~/Desktop/V3/v3/v2/Ishrak/pde_d_Poisson_2D_v3.py in load(self, name)
131
132 def load(self,name):
--> 133 restored_saved_model=keras.models.load_model('saved_model/')
134 return restored_saved_model
135 #Have to check
~/anaconda3/envs/tf2.1/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/save.py in load_model(filepath, custom_objects, compile)
148 if isinstance(filepath, six.string_types):
149 loader_impl.parse_saved_model(filepath)
--> 150 return saved_model_load.load(filepath, compile)
151
152 raise IOError(
~/anaconda3/envs/tf2.1/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py in load(path, compile)
87 # TODO(kathywu): Add saving/loading of optimizer, compiled losses and metrics.
88 # TODO(kathywu): Add code to load from objects that contain all endpoints
---> 89 model = tf_load.load_internal(path, loader_cls=KerasObjectLoader)
90
91 # pylint: disable=protected-access
~/anaconda3/envs/tf2.1/lib/python3.7/site-packages/tensorflow_core/python/saved_model/load.py in load_internal(export_dir, tags, loader_cls)
550 loader = loader_cls(object_graph_proto,
551 saved_model_proto,
--> 552 export_dir)
553 root = loader.get(0)
554 root.tensorflow_version = meta_graph_def.meta_info_def.tensorflow_version
~/anaconda3/envs/tf2.1/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py in __init__(self, *args, **kwargs)
117 def __init__(self, *args, **kwargs):
118 super(KerasObjectLoader, self).__init__(*args, **kwargs)
--> 119 self._finalize()
120
121 def _finalize(self):
~/anaconda3/envs/tf2.1/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py in _finalize(self)
137 for node in self._nodes:
138 if isinstance(node, RevivedNetwork):
--> 139 call_fn = node.keras_api.call_and_return_conditional_losses
140 if call_fn.input_signature is None:
141 inputs = infer_inputs_from_restored_call_function(call_fn)
AttributeError: '_UserObject' object has no attribute 'call_and_return_conditional_losses'
How do I save and load a TensorFlow model in this scenario?