From 3a6245112e4e2bd3010ef568dbd937d7b3303346 Mon Sep 17 00:00:00 2001
From: Ansh Sharma <60016461+DaemonOnCode@users.noreply.github.com>
Date: Tue, 24 Aug 2021 09:30:31 +0530
Subject: [PATCH] :ambulance: Fix memory growth amongst multiple GPU

---
 app.py               | 2 +-
 detect.py            | 2 +-
 detect_video.py      | 2 +-
 yolov3_tf2/models.py | 4 ++--
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/app.py b/app.py
index c317a09..c2915e5 100644
--- a/app.py
+++ b/app.py
@@ -22,7 +22,7 @@
 # load in weights and classes
 physical_devices = tf.config.experimental.list_physical_devices('GPU')
 if len(physical_devices) > 0:
-    tf.config.experimental.set_memory_growth(physical_devices[0], True)
+    tf.config.set_visible_devices(physical_devices[0:1], 'GPU')
 
 if tiny:
     yolo = YoloV3Tiny(classes=num_classes)
diff --git a/detect.py b/detect.py
index d0ade74..b3f3f2b 100644
--- a/detect.py
+++ b/detect.py
@@ -23,7 +23,7 @@
 def main(_argv):
     physical_devices = tf.config.experimental.list_physical_devices('GPU')
     if len(physical_devices) > 0:
-        tf.config.experimental.set_memory_growth(physical_devices[0], True)
+        tf.config.set_visible_devices(physical_devices[0:1], 'GPU')
 
     if FLAGS.tiny:
         yolo = YoloV3Tiny(classes=FLAGS.num_classes)
diff --git a/detect_video.py b/detect_video.py
index f54de1f..553b383 100644
--- a/detect_video.py
+++ b/detect_video.py
@@ -25,7 +25,7 @@
 def main(_argv):
     physical_devices = tf.config.experimental.list_physical_devices('GPU')
     if len(physical_devices) > 0:
-        tf.config.experimental.set_memory_growth(physical_devices[0], True)
+        tf.config.set_visible_devices(physical_devices[0:1], 'GPU')
 
     if FLAGS.tiny:
         yolo = YoloV3Tiny(classes=FLAGS.num_classes)
diff --git a/yolov3_tf2/models.py b/yolov3_tf2/models.py
index aa6d815..13ef549 100644
--- a/yolov3_tf2/models.py
+++ b/yolov3_tf2/models.py
@@ -208,7 +208,7 @@ def YoloV3(size=None, channels=3, anchors=yolo_anchors,
            masks=yolo_anchor_masks, classes=80, training=False):
     physical_devices = tf.config.experimental.list_physical_devices('GPU')
     if len(physical_devices) > 0:
-        tf.config.experimental.set_memory_growth(physical_devices[0], True)
+        tf.config.set_visible_devices(physical_devices[0:1], 'GPU')
     x = inputs = Input([size, size, channels], name='input')
 
     x_36, x_61, x = Darknet(name='yolo_darknet')(x)
@@ -242,7 +242,7 @@ def YoloV3Tiny(size=None, channels=3, anchors=yolo_tiny_anchors,
                masks=yolo_tiny_anchor_masks, classes=80, training=False):
     physical_devices = tf.config.experimental.list_physical_devices('GPU')
     if len(physical_devices) > 0:
-        tf.config.experimental.set_memory_growth(physical_devices[0], True)
+        tf.config.set_visible_devices(physical_devices[0:1], 'GPU')
     x = inputs = Input([size, size, channels], name='input')
 
     x_8, x = DarknetTiny(name='yolo_darknet')(x)