Merge pull request #40093 from ROCmSoftwarePlatform:google_upstream_rocm_csb_fix_200529
PiperOrigin-RevId: 314768988 Change-Id: Ic622943ca011fdfe80bc5751d9a9aa9e80a52019
This commit is contained in:
commit
3f8294f6f9
@ -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",
|
||||||
|
@ -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
|
||||||
|
|
||||||
|
@ -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",
|
||||||
|
Loading…
x
Reference in New Issue
Block a user