def METHOD_NAME(self, module):
    batch_size = 2
    data_specs = module.get_data_spec(batch_size * 2)
    dataset = get_dataset(data_specs, batch_size)
    for batch in dataset.take(1):
        pass
    if isinstance(batch, tuple):
        batch = batch[0]
    model = module.get_model(
        build=True,
        compile=True,
        jit_compile=False,
        include_preprocessing=True,
    )
    model.fit(dataset, epochs=1, steps_per_epoch=1)
    temp_filepath = os.path.join(
        self.get_temp_dir(), f"{module.__name__}.keras"
    )
    model.save(temp_filepath, save_format="keras_v3")
    with tf.keras.utils.custom_object_scope(module.get_custom_objects()):
        new_model = tf.keras.models.load_model(temp_filepath)
    # Test model weights
    self.assertIs(new_model.__class__, model.__class__)
    self.assertEqual(len(model.get_weights()), len(new_model.get_weights()))
    for w1, w2 in zip(model.get_weights(), new_model.get_weights()):
        if w1.dtype == "object":
            self.assertEqual(str(w1), str(w2))
        else:
            self.assertAllClose(w1, w2, atol=1e-6)
    # Test forward pass
    self.assertAllClose(new_model(batch), model(batch), atol=1e-6)
    # Test optimizer state
    if hasattr(model, "optimizer"):
        self.assertEqual(
            len(model.optimizer.variables()),
            len(new_model.optimizer.variables()),
        )
        for v1, v2 in zip(
            model.optimizer.variables(), new_model.optimizer.variables()
        ):
            self.assertAllClose(v1.numpy(), v2.numpy(), atol=1e-6)
    # Test training still works
    new_model.fit(dataset, epochs=1, steps_per_epoch=1)