| Geant4 Cross Reference |
1 import os
2 from dataclasses import dataclass
3
4 import tensorflow as tf
5
6
7 @dataclass
8 class GPULimiter:
9 """
10 Class responsible to set the limits of pos
11 instance of logical device per physical de
12
13 Attributes:
14 _gpu_ids: A string representing visibl
15 be separated by commas (no spaces)
16 _max_gpu_memory_allocation: An integer
17
18 """
19 _gpu_ids: str
20 _max_gpu_memory_allocation: int
21
22 def __call__(self):
23 os.environ["CUDA_VISIBLE_DEVICES"] = f
24 gpus = tf.config.list_physical_devices
25 if gpus:
26 # Restrict TensorFlow to only allo
27 try:
28 for gpu in gpus:
29 tf.config.set_logical_devi
30 gpu,
31 [tf.config.LogicalDevi
32 logical_gpus = tf.config.list_
33 print(len(gpus), "Physical GPU
34 except RuntimeError as e:
35 # Virtual devices must be set
36 print(e)