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", name = "gradient_checkpoint_test",
srcs = ["gradient_checkpoint_test.py"], srcs = ["gradient_checkpoint_test.py"],
python_version = "PY3", python_version = "PY3",
tags = ["no_rocm"],
deps = [ deps = [
"//tensorflow:tensorflow_py", "//tensorflow:tensorflow_py",
"//tensorflow/python:extra_py_tests_deps", "//tensorflow/python:extra_py_tests_deps",

View File

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

View File

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