title: Memory leaking after `tf.keras.Model.fit` is called and training doesn't start
body: <p>I'm using my yolo <a href="https://github.com/alternativebug/yolo-tf2" rel="nofollow noreferrer">implementation</a> which used to work fine on tensorflow versions prior to 2.5. I tried recently training yolo3 on a small dataset (which uses <code>tf.keras.Model.fit</code>). Here's a colab <a href="https://colab.research.google.com/drive/18jCTQajjgBO2bmKekaI5frrluJ2CX4NN?usp=sharing" rel="nofollow noreferrer">notebook</a> which you can use to reproduce the issue. Shortly after <code>model.fit</code> is called, the messages below keep repeating:</p>
  <pre><code>/usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  </code></pre>
  <p>and</p>
  <pre><code>INFO:tensorflow:Assets written to: ram://eefa3127-ad7d-4445-a186-75fd8f0b81e1/assets
  </code></pre>
  <p>Then memory usage keeps growing for no apparent reason and eventually a memory crash occurs. (which doesn't happen in earlier tensorflow versions &lt;= 2.5). You can verify so using this other <a href="https://colab.research.google.com/drive/1a3RAhVA3pCTQj2FyGY2df7vum27pj69a?usp=sharing" rel="nofollow noreferrer">notebook</a> which uses tensorflow 2.5 instead, things should go perfectly fine and training goes as expected. I also tried installing tensorflow 2.8 instead of colab's default version (2.7) and the issue persists.</p>
  <p>Here's the output containing problems (tensorflow &gt; 2.5):</p>
  <pre><code>2022-02-07 05:52:00,476 yolo_tf2.utils.common.activate_gpu +325: INFO     [260] GPU activated
  2022-02-07 05:52:00,477 yolo_tf2.utils.common.train +468: INFO     [260] Starting training ...
  2022-02-07 05:52:04,293 yolo_tf2.utils.common.create_models +447: INFO     [260] Training and inference models created
  2022-02-07 05:52:04,295 yolo_tf2.utils.common.wrapper +64: INFO     [260] create_models execution time: 3.8118433569999866 seconds
  2022-02-07 05:52:04,301 yolo_tf2.utils.common.create_new_dataset +366: INFO     [260] Generating new dataset ...
  2022-02-07 05:52:07,014 yolo_tf2.utils.common.adjust_non_voc_csv +184: INFO     [260] Adjustment from existing received 10107 labels containing 16 classes
  2022-02-07 05:52:07,022 yolo_tf2.utils.common.adjust_non_voc_csv +187: INFO     [260] Added prefix to images: /content/yolo-data/images
  Parsed labels:
  Car               3153
  Pedestrian        1418
  Palm Tree         1379
  Traffic Lights    1269
  Street Sign       1109
  Street Lamp        995
  Road Block         363
  Flag               124
  Trash Can           90
  Minivan             68
  Fire Hydrant        52
  Bus                 43
  Pickup Truck        20
  Bicycle             17
  Delivery Truck       4
  Motorcycle           3
  Name: object_name, dtype: int64
  2022-02-07 05:52:09,513 yolo_tf2.utils.common.save_fig +33: INFO     [260] Saved figure /content/output/plots/Relative width and height for 10107 boxes..png
  /usr/local/lib/python3.7/dist-packages/yolo_tf2/utils/dataset_handlers.py:209: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray
    groups = np.array(data.groupby('image_path'))
  Processing beverly_hills_train.tfrecord
  Building example: 406/411 ... Beverly_hills184.jpg 99% completed2022-02-07 05:52:12,922 yolo_tf2.utils.common.save_tfr +227: INFO     [260] Saved training TFRecord: /content/data/tfrecords/beverly_hills_train.tfrecord
  Building example: 411/411 ... Beverly_hills365.jpg 100% completed
  Processing beverly_hills_test.tfrecord
  Building example: 31/46 ... Beverly_hills335.jpg 67% completed2022-02-07 05:52:13,175 yolo_tf2.utils.common.save_tfr +229: INFO     [260] Saved validation TFRecord: /content/data/tfrecords/beverly_hills_test.tfrecord
  2022-02-07 05:52:13,271 yolo_tf2.utils.common.read_tfr +263: INFO     [260] Read TFRecord: /content/data/tfrecords/beverly_hills_train.tfrecord
  Building example: 46/46 ... Beverly_hills186.jpg 100% completed
  2022-02-07 05:52:18,892 yolo_tf2.utils.common.read_tfr +263: INFO     [260] Read TFRecord: /content/data/tfrecords/beverly_hills_test.tfrecord
  /usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  2022-02-07 05:52:50.575910: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
  INFO:tensorflow:Assets written to: ram://eefa3127-ad7d-4445-a186-75fd8f0b81e1/assets
  /usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  INFO:tensorflow:Assets written to: ram://cbe6d5a4-5322-494b-ba91-3fd34131cdd9/assets
  /usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  INFO:tensorflow:Assets written to: ram://f15f3f25-9adb-4eb0-aa0d-83fa874bc74e/assets
  /usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  INFO:tensorflow:Assets written to: ram://86dd6f5f-4416-4465-99c0-928fd88e8a93/assets
  /usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  INFO:tensorflow:Assets written to: ram://ca08220f-cabc-4017-96d3-383557342388/assets
  /usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  INFO:tensorflow:Assets written to: ram://0f634207-e822-4d6c-a805-3cfeab37532f/assets
  /usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  INFO:tensorflow:Assets written to: ram://a971d021-3da4-402a-a004-4ae4aa67148a/assets
  /usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  INFO:tensorflow:Assets written to: ram://31d72fdf-1ce6-4131-a7e6-f6444747e9c9/assets
  /usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  INFO:tensorflow:Assets written to: ram://dac323b6-591a-481c-bbe6-85bb82bef38c/assets
  /usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  INFO:tensorflow:Assets written to: ram://99b029f7-11d1-40f2-b459-fd1d8dca5ba1/assets
  /usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  INFO:tensorflow:Assets written to: ram://210489fb-0895-4769-8be3-effd01d92695/assets
  /usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    layer_config = serialize_layer_fn(layer)
  </code></pre>
  <p>Here's the output without the problem (tensorflow 2.5):</p>
  <pre><code>2022-02-07 06:09:53.125735: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
  2022-02-07 06:09:55.370728: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcuda.so.1
  2022-02-07 06:09:55,387 yolo_tf2.utils.common.train +468: INFO     [269] Starting training ...
  2022-02-07 06:09:55.387211: E tensorflow/stream_executor/cuda/cuda_driver.cc:328] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
  2022-02-07 06:09:55.387252: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (de0312867ce7): /proc/driver/nvidia/version does not exist
  2022-02-07 06:09:55.427963: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
  To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
  2022-02-07 06:10:00,078 yolo_tf2.utils.common.create_models +447: INFO     [269] Training and inference models created
  2022-02-07 06:10:00,080 yolo_tf2.utils.common.wrapper +64: INFO     [269] create_models execution time: 4.689235652999997 seconds
  2022-02-07 06:10:00,081 yolo_tf2.utils.common.create_new_dataset +366: INFO     [269] Generating new dataset ...
  2022-02-07 06:10:02,572 yolo_tf2.utils.common.adjust_non_voc_csv +184: INFO     [269] Adjustment from existing received 10107 labels containing 16 classes
  2022-02-07 06:10:02,574 yolo_tf2.utils.common.adjust_non_voc_csv +187: INFO     [269] Added prefix to images: /content/yolo-data/images
  Parsed labels:
  Car               3153
  Pedestrian        1418
  Palm Tree         1379
  Traffic Lights    1269
  Street Sign       1109
  Street Lamp        995
  Road Block         363
  Flag               124
  Trash Can           90
  Minivan             68
  Fire Hydrant        52
  Bus                 43
  Pickup Truck        20
  Bicycle             17
  Delivery Truck       4
  Motorcycle           3
  Name: object_name, dtype: int64
  2022-02-07 06:10:04,900 yolo_tf2.utils.common.save_fig +33: INFO     [269] Saved figure /content/output/plots/Relative width and height for 10107 boxes..png
  /usr/local/lib/python3.7/dist-packages/yolo_tf2/utils/dataset_handlers.py:209: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray
    groups = np.array(data.groupby('image_path'))
  Processing beverly_hills_train.tfrecord
  Building example: 392/411 ... Beverly_hills294.jpg 95% completed2022-02-07 06:10:10,341 yolo_tf2.utils.common.save_tfr +227: INFO     [269] Saved training TFRecord: /content/data/tfrecords/beverly_hills_train.tfrecord
  Building example: 411/411 ... Beverly_hills94.jpg 100% completed
  Processing beverly_hills_test.tfrecord
  Building example: 25/46 ... Beverly_hills334.jpg 54% completed2022-02-07 06:10:10,730 yolo_tf2.utils.common.save_tfr +229: INFO     [269] Saved validation TFRecord: /content/data/tfrecords/beverly_hills_test.tfrecord
  Building example: 46/46 ... Beverly_hills251.jpg 100% completed
  2022-02-07 06:10:10,843 yolo_tf2.utils.common.read_tfr +263: INFO     [269] Read TFRecord: /content/data/tfrecords/beverly_hills_train.tfrecord
  2022-02-07 06:10:15,264 yolo_tf2.utils.common.read_tfr +263: INFO     [269] Read TFRecord: /content/data/tfrecords/beverly_hills_test.tfrecord
  2022-02-07 06:10:15.676352: I tensorflow/core/profiler/lib/profiler_session.cc:126] Profiler session initializing.
  2022-02-07 06:10:15.676423: I tensorflow/core/profiler/lib/profiler_session.cc:141] Profiler session started.
  2022-02-07 06:10:15.701051: I tensorflow/core/profiler/lib/profiler_session.cc:159] Profiler session tear down.
  /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/utils/generic_utils.py:497: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.
    category=CustomMaskWarning)
  2022-02-07 06:10:17.064324: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:176] None of the MLIR Optimization Passes are enabled (registered 2)
  2022-02-07 06:10:17.081408: I tensorflow/core/platform/profile_utils/cpu_utils.cc:114] CPU Frequency: 2199995000 Hz
  Epoch 1/100
        1/Unknown - 40s 40s/step - loss: 7333.2617 - layer_205_lambda_loss: 403.8862 - layer_230_lambda_loss: 1509.9465 - layer_255_lambda_loss: 5407.68902022-02-07 06:10:59.974130: I tensorflow/core/profiler/lib/profiler_session.cc:126] Profiler session initializing.
  2022-02-07 06:10:59.974196: I tensorflow/core/profiler/lib/profiler_session.cc:141] Profiler session started.
        2/Unknown - 50s 11s/step - loss: 7819.7124 - layer_205_lambda_loss: 697.4546 - layer_230_lambda_loss: 1647.7856 - layer_255_lambda_loss: 5462.71582022-02-07 06:11:10.059899: I tensorflow/core/profiler/lib/profiler_session.cc:66] Profiler session collecting data.
  2022-02-07 06:11:10.088821: I tensorflow/core/profiler/lib/profiler_session.cc:159] Profiler session tear down.
  2022-02-07 06:11:10.133747: I tensorflow/core/profiler/rpc/client/save_profile.cc:137] Creating directory: /content/data/tfrecords/train/plugins/profile/2022_02_07_06_11_10
  2022-02-07 06:11:10.157875: I tensorflow/core/profiler/rpc/client/save_profile.cc:143] Dumped gzipped tool data for trace.json.gz to /content/data/tfrecords/train/plugins/profile/2022_02_07_06_11_10/de0312867ce7.trace.json.gz
  2022-02-07 06:11:10.189438: I tensorflow/core/profiler/rpc/client/save_profile.cc:137] Creating directory: /content/data/tfrecords/train/plugins/profile/2022_02_07_06_11_10
  2022-02-07 06:11:10.189678: I tensorflow/core/profiler/rpc/client/save_profile.cc:143] Dumped gzipped tool data for memory_profile.json.gz to /content/data/tfrecords/train/plugins/profile/2022_02_07_06_11_10/de0312867ce7.memory_profile.json.gz
  2022-02-07 06:11:10.192796: I tensorflow/core/profiler/rpc/client/capture_profile.cc:251] Creating directory: /content/data/tfrecords/train/plugins/profile/2022_02_07_06_11_10Dumped tool data for xplane.pb to /content/data/tfrecords/train/plugins/profile/2022_02_07_06_11_10/de0312867ce7.xplane.pb
  Dumped tool data for overview_page.pb to /content/data/tfrecords/train/plugins/profile/2022_02_07_06_11_10/de0312867ce7.overview_page.pb
  Dumped tool data for input_pipeline.pb to /content/data/tfrecords/train/plugins/profile/2022_02_07_06_11_10/de0312867ce7.input_pipeline.pb
  Dumped tool data for tensorflow_stats.pb to /content/data/tfrecords/train/plugins/profile/2022_02_07_06_11_10/de0312867ce7.tensorflow_stats.pb
  Dumped tool data for kernel_stats.pb to /content/data/tfrecords/train/plugins/profile/2022_02_07_06_11_10/de0312867ce7.kernel_stats.pb
  
       15/Unknown - 181s 10s/step - loss: 3493.4009 - layer_205_lambda_loss: 232.7220 - layer_230_lambda_loss: 629.6332 - layer_255_lambda_loss: 2618.9722
  </code></pre>
  <p>I also tried (same results):</p>
  <ul>
  <li>python versions: 3.8, 3.9, 3.10</li>
  <li>Ubuntu 18, and OSX</li>
  <li>tensorflow versions: 2.7, 2.8, 2.9.0-dev20220203</li>
  </ul>