Yolo v3 Loss Function fails when reloaded during load_model - tensorflow

Question

I've been trying to save a Yolo v3 model and then Load it back from from an h5 file.
When saving I use the checkpoint (ModelCheckpoint) to save the model (with the parameter save_weights_only
set to False in order to save the WHOLE model).
However, when I tried to recover the same model by using the keras load_model function, I initially get a yolo_head function not found error.
I then tried to add the function as a parameter to the load function as in:
{"yolo_head":yolo_head}
Now, the issue becomes: "TypeError: list indices must be integers or slices, not list" because somehow, there's an error in the loss function (yolo_loss, line 444) when loaded dynamically.
Apparently, the binary code of the loss function is hard copied into the h5 file.
My question is this:
Is there a better/simpler YOLO loss function that I can use THAT DOES NOT refer to other functions or can be easily reloaded?
Thanks in advance,
EDIT 1: Additional Code Snippets,
Keras Checkpoint Callback definition:
checkpoint = ModelCheckpoint(
os.path.join(log_dir, "checkpoint.h5"),
monitor="val_loss",
save_weights_only=False,
save_best_only=True,
period=1,
)
Checkpoint added to model training:
history = model.fit_generator(
data_generator_wrapper(
lines[:num_train], batch_size, input_shape, anchors, num_classes
),
steps_per_epoch=max(1, num_train // batch_size),
validation_data=data_generator_wrapper(
lines[num_train:], batch_size, input_shape, anchors, num_classes
),
validation_steps=max(1, num_val // batch_size),
epochs=epoch1,
initial_epoch=0,
callbacks=[logging, checkpoint],
)
Trying to load the same file 'checkpoint.h5' after pre-training ended:
weights_path = os.path.join(log_dir, "checkpoint.h5")
model = load_model(weights_path, {"yolo_head":yolo_head, "tf":tf, "box_iou":box_iou,'<lambda>': lambda y_true, y_pred: y_pred})
Here is the error stack trace:
File "2_Training/Train_YOLO.py", line 206, in
model = load_model(weights_path, {"yolo_head":yolo_head, "tf":tf,
"box_iou":box_iou,'': lambda y_true, y_pred: y_pred})
File "/Users/nkwedi/.pyenv/versions/3.7.5/lib/python3.7/site-packages/keras/engine/saving.py", line 419, in load_model
model = _deserialize_model(f, custom_objects, compile)
File "/Users/nkwedi/.pyenv/versions/3.7.5/lib/python3.7/site-packages/keras/engine/saving.py", line 225, in _deserialize_model
model = model_from_config(model_config, custom_objects=custom_objects)
File "/Users/nkwedi/.pyenv/versions/3.7.5/lib/python3.7/site-packages/keras/engine/saving.py", line 458, in model_from_config
return deserialize(config, custom_objects=custom_objects)
File "/Users/nkwedi/.pyenv/versions/3.7.5/lib/python3.7/site-packages/keras/layers/init.py", line 55, in deserialize
printable_module_name='layer')
File "/Users/nkwedi/.pyenv/versions/3.7.5/lib/python3.7/site-packages/keras/utils/generic_utils.py", line 145, in deserialize_keras_object
list(custom_objects.items())))
File "/Users/nkwedi/.pyenv/versions/3.7.5/lib/python3.7/site-packages/keras/engine/network.py", line 1032, in from_config
process_node(layer, node_data)
File "/Users/nkwedi/.pyenv/versions/3.7.5/lib/python3.7/site-packages/keras/engine/network.py", line 991, in process_node
layer(unpack_singleton(input_tensors), **kwargs)
File "/Users/nkwedi/.pyenv/versions/3.7.5/lib/python3.7/site-packages/keras/engine/base_layer.py", line 457, in call
output = self.call(inputs, **kwargs)
File "/Users/nkwedi/.pyenv/versions/3.7.5/lib/python3.7/site-packages/keras/layers/core.py", line 687, in call
return self.function(inputs, **arguments)
File "/Users/nkwedi/Documents/MyProjects/Eroscope/EyeDetectionYOLO/2_Training/src/keras_yolo3/yolo3/model.py", line 444, in yolo_loss
anchors[anchor_mask[l]],
TypeError: list indices must be integers or slices, not list

Answer 1

The Solution for me was to use a Cloud Based Training Platform like Google Collab.
Here's a link to a workable Collab Notebook with GPU enabled:
YOLO v3 Google Collab Tutorial