Merge pull request #40093 from ROCmSoftwarePlatform:google_upstream_rocm_csb_fix_200529

PiperOrigin-RevId: 314768988
Change-Id: Ic622943ca011fdfe80bc5751d9a9aa9e80a52019
This commit is contained in:
TensorFlower Gardener 2020-06-04 11:38:05 -07:00
commit 3f8294f6f9
3 changed files with 3 additions and 1 deletions

View File

@ -75,6 +75,7 @@ cuda_py_test(
name = "gradient_checkpoint_test",
srcs = ["gradient_checkpoint_test.py"],
python_version = "PY3",
tags = ["no_rocm"],
deps = [
"//tensorflow:tensorflow_py",
"//tensorflow/python:extra_py_tests_deps",

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=1152)])
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
return True
return False

View File

@ -105,6 +105,7 @@ cuda_py_test(
size = "medium",
srcs = ["adam_test.py"],
shard_count = 4,
tags = ["no_rocm"],
deps = [
":optimizer_v2",
"//tensorflow/python:client_testlib",