Error "Unknown layer: 'LayerScale'" with the ConvNeXtXLarge model when testing it with an X input

Hello everyone I have trained the pre-trained ConvNeXtXLarge model of Keras Application on my image dataset. When I try to load my model and perform a prediction on an input X I get the following error:

Traceback (most recent call last):
  File "/multiverse/storage/lattari/Prj/postdoc/Courses/AN2DL_2023/Competition1_running_dir/worker_gpu4_dir/tmp/codalab/tmpWeFCGo/run/program/", line 130, in 
    M = model(submission_dir)
  File "/multiverse/storage/lattari/Prj/postdoc/Courses/AN2DL_2023/Competition1_running_dir/worker_gpu4_dir/tmp/codalab/tmpWeFCGo/run/input/res/", line 6, in __init__
    self.model = tf.keras.models.load_model(os.path.join(path, 'SubmissionModel', 'my_model.keras'))
  File "/usr/local/lib/python3.11/dist-packages/keras/src/saving/", line 262, in load_model
    return legacy_sm_saving_lib.load_model(
  File "/usr/local/lib/python3.11/dist-packages/keras/src/utils/", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/local/lib/python3.11/dist-packages/keras/src/saving/legacy/", line 365, in class_and_config_for_serialized_keras_object
    raise ValueError(
ValueError: Unknown layer: 'LayerScale'. Please ensure you are using a `keras.utils.custom_object_scope` and that this object is included in the scope. See for details.

Does anyone have any idea why I am getting this error?

the model you are trying to load contains a custom layer named ‘LayerScale’ that is not recognized.

import tensorflow as tf

Define LayerScale class

class LayerScale(tf.keras.layers.Layer):
# Your implementation here

Wrap model loading within a custom object scope

with keras.utils.custom_object_scope({‘LayerScale’: LayerScale}):
model = tf.keras.models.load_model(‘path_to_your_model’)
try this.