Found the solution, I should use TF_LoadPluggableDeviceLibrary API, but it crashes after loading the plugin 2022-08-17 08:16:29.397578: I tensorflow/cc/saved_model/reader.cc:43] Reading SavedModel 2022-08-17 08:16:29.401477: I tensorflow/cc/saved_model/reader.cc:81] Reading meta graph with tags { serve } 2022-08-17 08:16:29.401500: I tensorflow/cc/saved_model/reader.cc:122] Reading SavedModel debug info (if present) Metal device set to: Apple M1 systemMemory: 16.00 GB maxCacheSize: 5.33 GB 2022-08-17 08:16:29.411250: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:306] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support. 2022-08-17 08:16:29.411398: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:272] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>) 2022-08-17 08:16:29.422914: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:354] MLIR V1 optimization pass is not enabled 2022-08-17 08:16:29.427004: I tensorflow/cc/saved_model/loader.cc:227] Restoring SavedModel bundle. 2022-08-17 08:16:29.429435: W tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz 2022-08-17 08:16:29.440099: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:113] Plugin optimizer for device_type GPU is enabled. zsh: segmentation fault

Hello I am encountering the same issue in my Mac M1. I have a saved model for prediction only. It predicts one batch after another. Tensorflow v2.10.0 tensorflow-metal v0.6.0 When using CPU, only consume ~400MB memory. When starting using GPU via tf-metal, RAM consumption keeps increasing and never stops as long as prediction does not stop. I have 16G RAM in total, and it can easily run out of my RAM.