Python ResNet50: model.save() NotImplementedError

My goal is to save (and then load) a resent model. I've followed this tutorial and I've ended up with a model that learns, but when I try to save it, it gives an error.

I've found this similar stackoverflow issue, but for the life of me I could not figure out how to solve it.

Another thing I looked at was this article from Keras.io, but I'm using Sequential() model and not some custom one. I'm not sure where should this get_config function should be.

import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.python.keras.layers import Dense, Flatten
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam

import matplotlib.pyplot as plt


DATASET_PATH = "/XX/dataset"
CLASS_NAMES = ["0", "1", "2", "3", "4"]

img_height,img_width=180,180
batch_size=32

train_ds = tf.keras.preprocessing.image_dataset_from_directory(
  DATASET_PATH,
  validation_split=0.2,
  subset="training",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size)

val_ds = tf.keras.preprocessing.image_dataset_from_directory(
  DATASET_PATH,
  validation_split=0.2,
  subset="validation",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size)

resnet_model = Sequential()

pretrained_model= tf.keras.applications.ResNet50(include_top=False,
       input_shape=(180,180,3),
       pooling='avg',classes=5,
       weights='imagenet')

for layer in pretrained_model.layers:
    layer.trainable=False

resnet_model.add(pretrained_model)
resnet_model.add(Flatten())
resnet_model.add(Dense(512, activation='relu'))
resnet_model.add(Dense(5, activation='softmax'))

resnet_model.summary()

resnet_model.compile(optimizer=Adam(lr=0.001),loss='sparse_categorical_crossentropy',metrics=['accuracy'])

epochs=1
history = resnet_model.fit(
  train_ds,
  validation_data=val_ds,
  epochs=epochs
)

resnet_model.save("/XX/test.h5", save_format="h5")

And the error:

NotImplementedError: 
Layer ModuleWrapper has arguments ['self', 'module', 'method_name']
in `__init__` and therefore must override `get_config()`.

Example:

class CustomLayer(keras.layers.Layer):
    def __init__(self, arg1, arg2):
        super().__init__()
        self.arg1 = arg1
        self.arg2 = arg2

    def get_config(self):
        config = super().get_config()
        config.update({
            "arg1": self.arg1,
            "arg2": self.arg2,
        })
        return config

The problem is with this line

from tensorflow.python.keras.layers import Dense, Flatten

If you replace it to this it should solve your problem

from tensorflow.keras.layers import Dense, Flatten