When running BERT on GPU: Resource exhausted: failed to allocate memory

[Currycurrycurry]

File "/root/softwares/anaconda3/envs/tf115/lib/python3.6/site-packages/tensorflow/python/client/session.py",
line 1453, in _call_tf_sessionrun
run_metadata)
tensorflow.python.framework.errors_impl.ResourceExhaustedError: 2 root error(s) found.
(0) Resource exhausted: failed to allocate memory
[[{{node bert/encoder/block_0/feedforward_1/dense_1/scalar_mul/parallel_4/mul}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom
to RunOptions for current allocation info. This isn't available when running in Eager mode.

     [[reshape/parallel_0/Reshape/_1265]]                                                                 

Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom
to RunOptions for current allocation info. This isn't available when running in Eager mode.

(1) Resource exhausted: failed to allocate memory
[[{{node bert/encoder/block_0/feedforward_1/dense_1/scalar_mul/parallel_4/mul}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom
to RunOptions for current allocation info. This isn't available when running in Eager mode.

0 successful operations.
0 derived errors ignored.

[Currycurrycurry]

Anyone can help? @cghawthorne @crccw @lucidrains @ronw