Merge pull request #39914 from ROCmSoftwarePlatform:google_upstream_rocm_csb_fix_200527

PiperOrigin-RevId: 313520606
Change-Id: I0c174a0a8a21a05de370ae8f6afa38a0cb5008e0
This commit is contained in:
TensorFlower Gardener 2020-05-27 21:24:01 -07:00
commit f150cad3cb

View File

@ -75,7 +75,7 @@ def _limit_gpu_memory():
if gpus:
tf.config.experimental.set_virtual_device_configuration(
gpus[0],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1152)])
return True
return False