Saving the model weights in TF 2

[danielstafford2001]

When completing this project using TF2 many functions in base_model can't work due to TF2 not using Sessions. In evaluate.py the function restore_session() is called, what changes should be made to the code to do this in TF2.

I have tried
checkpoint_path = "training_3/cp.ckpt"#defining path where weights will be saved
checkpoint_dir = os.path.dirname(checkpoint_path)

cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
save_weights_only=True,
verbose=1)#defining a function that is called in my model.fit by using: callbacks=[cp_callback])

I can then call load_weights but I am struggling with doing this within the main() in evaluate.py
Any help with this issue would be great, thanks.