SageMaker script mode inference model init needed

Hi, first post here, so please bear with me if it's not exactly according to the rules, I'll try to provide more info if needed.
I have successfully trained a keras model which in turn uses a tensorflow hub model (specifically bert-tensorflow) as the first layer with SageMaker using script mode.
To make this model work, in the training script I have to call these 4 lines:
`sess.run(tf.compat.v1.local_variables_initializer())
sess.run(tf.compat.v1.global_variables_initializer())
sess.run(tf.compat.v1.tables_initializer())
tf.keras.backend.set_session(sess)`
as they are required to initialise the tensorflow hub bert model.
However, when the code is deployed to an endpoint, I'm getting the following error:
`ModelError: An error occurred (ModelError) when calling the InvokeEndpoint operation: Received client error (400) from model with message "{ "error": "Error while reading resource variable bert_layer_module/bert/encoder/layer_5/attention/self/key/kernel from Container: localhost. This could mean that the variable was uninitialized. Not found: Resource localhost/bert_layer_module/bert/encoder/layer_5/attention/self/key/kernel/N10tensorflow3VarE does not exist.\n\t [[{{node bert_layer/bert_layer_module_apply_tokens/bert/encoder/layer_5/attention/self/key/MatMul/ReadVariableOp}} = ReadVariableOp[_output_shapes=[[768,768]], dtype=DT_FLOAT, _device=\"/job:localhost/replica:0/task:0/device:CPU:0\"](bert_layer_module/bert/encoder/layer_5/attention/self/key/kernel)]]" }". See https://eu-west-2.console.aws.amazon.com/cloudwatch/home?region=eu-west-2#logEventViewer:group=/aws/sagemaker/Endpoints/Test in account 124774543455 for more information.`
which basically means that these lines were not run on the endpoint machine and therefore the variables have not been initialised and tensorflow can't find them.
How is it possible for me to make these lines run on the endpoint machine before doing prediction, as they must be run once before doing inference?
I tried having a separate inference.py script which runs these lines and feed it as an entry_point when creating the tensorflow model in the following way:
`tf_model = TensorFlowModel(model_data=model,
role=sagemaker.get_execution_role(),
entry_point='inference.py',
source_dir='.',
py_version='py3',
env={'SAGEMAKER_REQUIREMENTS': 'requirements.txt'})`
however, this didn't help.
Moreover, when doing inference locally, I don't run these lines but instead run the following line:
`tf.keras.backend.manual_variable_initialization(True)
model = tf.keras.models.load_model(checkpoint_file_name, custom_objects={'BertLayer': BertLayer})`
because the `load_model` function runs these lines instead, and if I were to run those lines it would reset the loaded variables and the model would lose its training.
In this case, since there's no load_model function with how script mode works, I don't quite understand how to do it properly.

I will really appreciate any help you could provide me with, as I've been struggling with this issue for quite a while now and the documentation is very confusing because most of it is outdated and doesn't apply to script mode.

Thank you