From 3ea1568ad34661237894c4e04ad22e467b0b218a Mon Sep 17 00:00:00 2001
From: Jordon Brooks <jordonbc@hotmail.co.uk>
Date: Thu, 17 Aug 2023 01:57:43 +0100
Subject: [PATCH] Added GPU growth

---
 train_model.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/train_model.py b/train_model.py
index dd93f80..b122989 100644
--- a/train_model.py
+++ b/train_model.py
@@ -14,8 +14,18 @@ from featureExtraction import psnr
 
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
 
+import gc
 import tensorflow as tf
-from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, Callback
+
+gpus = tf.config.experimental.list_physical_devices('GPU')
+if gpus:
+    try:
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+    except RuntimeError as e:
+        print(e)
+
 
 from video_compression_model import VideoCompressionModel, data_generator