Update app.py

app.py

@@ -3,6 +3,16 @@ import spaces
 import torch
 import keras
 import keras_nlp
+import tensorflow as tf
+
+gpus = tf.config.experimental.list_physical_devices('GPU')
+if gpus:
+    try:
+        # Restrict TensorFlow to only allocate memory as needed
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+    except RuntimeError as e:
+        print(e)
 
 keras.utils.set_random_seed(42)