SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Use stable functions (#9369)

Browse Source
This commit is contained in:
Julien Plu
2021-01-05 09:58:26 +01:00
committed by GitHub
parent 4aa8f6ad99
commit 4225740a7b
5 changed files with 12 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/test_modeling_tf_common.py
View File

@@ -51,10 +51,10 @@ if is_tf_available():
for gpu in gpus:
# Restrict TensorFlow to only allocate x GB of memory on the GPUs
try:
tf.config.experimental.set_virtual_device_configuration(
gpu, [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=_tf_gpu_memory_limit)]
tf.config.set_logical_device_configuration(
gpu, [tf.config.LogicalDeviceConfiguration(memory_limit=_tf_gpu_memory_limit)]
)
logical_gpus = tf.config.experimental.list_logical_devices("GPU")
logical_gpus = tf.config.list_logical_devices("GPU")
print("Logical GPUs", logical_gpus)
except RuntimeError as e:
# Virtual devices must be set before GPUs have been initialized