Properly save custom DRL Agent in TF 2

So I have a DDQN model, that looks like this:

class DDQN(keras.Model):
def __init__(self, n_actions, fc1_dims, fc2_dims):
super(DDQN, self).__init__()
self.dense1 = keras.layers.Dense(fc1_dims, activation=’relu’)
self.dense1.trainable = True
self.dense2 = keras.layers.Dense(fc2_dims, activation=’relu’)
self.dense1.trainable = True
self.V = keras.layers.Dense(1, activation=None) #Value stream layer
self.V.trainable = True
self.A = keras.layers.Dense(n_actions, activation=None) #Advantage stream layer
self.A.trainable = True
def call(self, state):
x = self.dense1(state)
x = self.dense2(x)
V = self.V(x)
A = self.A(x)
Q = (V + (A – tf.math.reduce_mean(A, axis=1, keepdims=True)))
return Q
def advantage(self, state):
x = self.dense1(state)
x = self.dense2(x)
A = self.A(x)
return A

I do call my model like this:

self.q_eval = DDQN(n_actions, fc1_dims, fc2_dims)

with tf.device(device):
self.q_eval.compile(optimizer=Adam(learning_rate=lr), loss=’mean_squared_error’)

​

now training and all works, results are great n stuff.

But that’s all pointless, if I can’t save the model, right?

Well, that’s where the trouble began.

​

def save_model(self):
self.q_eval.save(‘test’)

​

whenever I call this, I get this error:

​

Exception has occurred: ValueError

Model <__main__.DDQN object at 0x0000014483B8E948> cannot be saved because the input shapes have not been set. Usually, input shapes are automatically determined from calling `.fit()` or `.predict()`. To manually set the shapes, call `model.build(input_shape)`.

​

If I do model.build() n shit before saving, nothing changes. Now what would be the proper way to save a model like this?

Please excuse spelling mistakes or dumb sentences in the post and comments, since English is not my native language.