added '--allow-growth' flag to choose to limit GPU growth. This solves failed to create cublas handle: CUBLAS_STATUS_NOT_INITIALIZED issue if caused due to GPU memory limit.

Description

While executing training for object detection model in TF 2.x, the following error was thrown:

failed to create cublas handle: CUBLAS_STATUS_NOT_INITIALIZED

Issue #35029 already discusses the solution. Hence incorporated the same into model_main_tf2.py. Users can now opt to limit GPU memory growth by setting the --allow-growth flag while executing model_main_tf2.py.

Type of change

For a new feature or function, please create an issue first to discuss it
with us before submitting a pull request.

Note: Please delete options that are not relevant.

• TensorFlow 2 migration
• New feature (non-breaking change which adds functionality)

Checklist

• I have signed the Contributor License Agreement.
• I have read guidelines for pull request.
• My code follows the coding guidelines.
• I have performed a self code review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have made corresponding changes to the documentation. NA
• My changes generate no new warnings.
• I have added tests that prove my fix is effective or that my feature works. NA