Hi @Hapero!
I think that is an issue that has to do with the train_step function, when I reproduce it on my personal computer it rises me that OOM error too. For solving it try to decorate the train_step function like this:
#Decorator @tf.function for faster training in graph mode and avoiding OOM errors
@tf.function(experimental_relax_shapes=True)
def train_step(self, image_list, gt_boxes, gt_classes, optimizer, fine_tune_variables):
Enabling experimental_relax_shapes, may generate fewer graphs that are less specialized on input shapes and that would help the training loop to work with bigger datasets.
If the train_step funcion works correctly, try decorating the function you think that could potentially be responsible for the OOM with the experimental_relax_shapes option to True.
I wish it helps, tell me if you have any other problem. Keep it up!