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",
|
||||
srcs = ["gradient_checkpoint_test.py"],
|
||||
python_version = "PY3",
|
||||
tags = ["no_rocm"],
|
||||
deps = [
|
||||
"//tensorflow:tensorflow_py",
|
||||
"//tensorflow/python:extra_py_tests_deps",
|
||||
|
@ -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
|
||||
|
||||
|
@ -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",
|
||||
|
Loading…
x
Reference in New Issue
Block a user