From 9e65290e148db00ff27092897a34c72bb1bae1da Mon Sep 17 00:00:00 2001 From: Eli Mernit Date: Wed, 2 Oct 2024 17:18:46 -0400 Subject: [PATCH] Eli/update examples (#48) * add whisperx example and update flux volume path * add tf example --- gpu_acceleration/using_tensorflow/app.py | 27 ++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 gpu_acceleration/using_tensorflow/app.py diff --git a/gpu_acceleration/using_tensorflow/app.py b/gpu_acceleration/using_tensorflow/app.py new file mode 100644 index 0000000..7b16ba6 --- /dev/null +++ b/gpu_acceleration/using_tensorflow/app.py @@ -0,0 +1,27 @@ +from beam import Image, endpoint, env + +if env.is_remote(): + import tensorflow as tf + + +@endpoint( + name="tensorflow-gpu", + cpu=1, + memory="4Gi", + gpu="A10G", + # Make sure to use `tensorflow[and-cuda]` in order to access GPU resources + image=Image().add_python_packages(["tensorflow[and-cuda]"]), +) +def predict(): + # Show available GPUs + gpus = tf.config.list_physical_devices("GPU") + + try: + for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + except RuntimeError as e: + print(e) + + print("🚧 Is built with CUDA:", tf.test.is_built_with_cuda()) + print("🚧 Is GPU available:", tf.test.is_gpu_available()) + print("🚧 GPUs available:", tf.config.list_physical_devices("GPU"))