pat2vec.util.get_best_gpu
Functions
|
Selects the best available GPU based on free memory. |
- pat2vec.util.get_best_gpu.set_best_gpu(gpu_mem_threshold)[source]
Selects the best available GPU based on free memory.
This function checks for available CUDA-enabled GPUs and queries their free memory. If a GPU is found with free memory exceeding the specified threshold, it sets the CUDA_VISIBLE_DEVICES environment variable to that GPU’s index, effectively selecting it for subsequent processes. If no suitable GPU is found, it sets the variable to “-1” to force CPU-only mode.
- Parameters:
gpu_mem_threshold (
int) – The minimum amount of free memory (in MB) required to select a GPU.- Return type:
None