AWS Sagemaker edge failed to build due to Keras framework

Hi Team,

We have trained our Ml model using keras framework (version =2.11.0) which is used for prediction, and we want to run this model on an edge device (linux) so we made use of aws sagemaker edge manager but while compiling this model we are facing below error,

Error: localhost compiler-container-Primary[5139]: RuntimeError: InputConfiguration: Unable to load provided Keras model. Error: 'sample_weight_mode'

And as per the documentation (https://docs.aws.amazon.com/sagemaker/latest/dg/neo-supported-cloud.html), the compiler uses keras version 2.2.4 which does not support the features of 2.11, so we tried to use the 2.2.4 keras version to train our model but we were not able to achieve this as the tensorflow is updating the keras version.

So is there a possibility that we can compile the model using the latest keras version, or any other solution from your end which can help us in resolving this issue.

Thanks, Bala