memory error when runing mayo_learned_primal_dual

[hjahan58]

hello.
thanks for your attention.
I trying to train mayo_learned_primal_dual with some database.
I run program on gpu enabled env Quadro RTX 4000 with 8G memory .but i get OOM Error
I run program on cpu enabled env with 32G ram. but I get OOM error.

.
.
.
2021-06-07 17:42:42.878792: I tensorflow/core/common_runtime/bfc_allocator.cc:645] Sum Total of in-use chunks: 25.98GiB
2021-06-07 17:42:42.878807: I tensorflow/core/common_runtime/bfc_allocator.cc:647] Stats:
Limit: 33547444224
InUse: 27897504512
MaxInUse: 27897504512
NumAllocs: 536
MaxAllocSize: 6710886400

2021-06-07 17:42:42.878844: W tensorflow/core/common_runtime/bfc_allocator.cc:271] *****______**___
2021-06-07 17:42:43.028891: W tensorflow/core/framework/op_kernel.cc:1273] OP_REQUIRES failed at mkl_relu_op.cc:861 : Resource exhausted: OOM when allocating tensor with shape[1677721600] and type float on /job:localhost/replica:0/task:0/device:CPU:0 by allocator mklcpu
Traceback (most recent call last):
.
.
.