From ea5d80fa4a70eecfe1fb89b8c7407c5c4b194b95 Mon Sep 17 00:00:00 2001
From: Hunor Laczko <luni741@gmail.com>
Date: Fri, 30 Sep 2022 21:55:29 +0000
Subject: [PATCH] smpl memory limit

---
 custom_utils/smpl.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/custom_utils/smpl.py b/custom_utils/smpl.py
index 68a097e..5e7dc34 100644
--- a/custom_utils/smpl.py
+++ b/custom_utils/smpl.py
@@ -8,6 +8,18 @@ from math import *
 import numpy as np
 import tensorflow as tf
 
+gpus = tf.config.experimental.list_physical_devices('GPU')
+if gpus:
+    try:
+        # Currently, memory growth needs to be the same across GPUs
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
+        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
+    except RuntimeError as e:
+        # Memory growth must be set before GPUs have been initialized
+        print(e)
+
 # Left-Right
 eval_joints = {
         'All': [1, 2, 3, 4, 5, 7, 8, 9, 12, 15, 16, 17, 18, 19, 20, 21],