Out-of-memory problem submitting a tensorflow2 job on Google AI Platform Engine

Prerequisites

Please answer the following questions for yourself before submitting an issue.

• I am using the latest TensorFlow Model Garden release and TensorFlow 2.
• I am reporting the issue to the correct repository. (Model Garden official or research directory)
• I checked to make sure that this issue has not already been filed.

1. The entire URL of the file you are using

https://github.com/tensorflow/models/tree/master/research/object_detection/model_main_tf2.py

2. Describe the bug

I'm trying to submit a Tensorflow2 training job (fine tuning an object detection model) with gcloud on Google AI Platform Engine. My dataset is not big (raccoon dataset, which is 10M or so). I've tried many configurations but each time get the same error: "The replica master 0 ran out-of-memory and exited with a non-zero status of 9(SIGKILL)"

3. Steps to reproduce

docker install of OD

cloud.yml file content:
trainingInput:
runtimeVersion: "2.2"
pythonVersion: "3.7"
scaleTier: CUSTOM
masterType: large_model
workerCount: 5
workerType: large_model
parameterServerCount: 3
parameterServerType: large_model

executed command:
gcloud ai-platform jobs submit training OD_ssd_fpn_large \ --job-dir=gs://${MODEL_DIR} \ --package-path ./object_detection \ --module-name object_detection.model_main_tf2 \ --region us-east1 \ --config cloud.yml \ -- \ --model_dir=gs://${MODEL_DIR} \ --pipeline_config_path=gs://${PIPELINE_CONFIG_PATH}

4. Expected behavior

Job finishes with Success

5. Additional context

6. System information

Google VM on AI platform. As said above, last try involved large_model machine type