Get error "Non-OK-status: GpuLaunchKernel" while running save_model.py

[edwardchang0112-nctu]

After running

python3 save_model.py --weights ./data/custom.weights --output ./checkpoints/custom-416 --input_size 416 --model yolov4 

got the following error messages

Non-OK-status: GpuLaunchKernel(FillPhiloxRandomKernelLaunch<Distribution>, num_blocks, block_size, 0, d.stream(), key, counter, gen, data, size, dist) status: Internal: an illegal memory access was encountered
Fatal Python error: Aborted

Current thread 0x0000007fa037b010 (most recent call first):
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/gen_random_ops.py", line 634 in random_standard_normal
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/random_ops.py", line 95 in random_normal
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/dispatch.py", line 206 in wrapper
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/init_ops_v2.py", line 1073 in random_normal
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/init_ops_v2.py", line 420 in __call__
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/resource_variable_ops.py", line 1722 in _init_from_args
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/resource_variable_ops.py", line 1595 in __init__
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variables.py", line 264 in __call__
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variable_scope.py", line 2626 in default_variable_creator
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variables.py", line 199 in <lambda>
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variables.py", line 221 in _variable_v1_call
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variables.py", line 260 in __call__
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer_utils.py", line 139 in make_variable
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/training/tracking/base.py", line 815 in _add_variable_with_custom_getter
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py", line 655 in add_weight
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/layers/convolutional.py", line 204 in build
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py", line 2625 in _maybe_build
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py", line 878 in _infer_output_signature
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py", line 840 in _keras_tensor_symbolic_call
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py", line 1108 in _functional_construction_call
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py", line 970 in __call__
  File "/home/fiduciaedge/Desktop/yolov4-custom-functions-master/core/common.py", line 31 in convolutional
  File "/home/fiduciaedge/Desktop/yolov4-custom-functions-master/core/yolov4.py", line 125 in YOLOv4
  File "/home/fiduciaedge/Desktop/yolov4-custom-functions-master/core/yolov4.py", line 25 in YOLO
  File "save_model.py", line 20 in save_tf
  File "save_model.py", line 54 in main
  File "/usr/local/lib/python3.6/dist-packages/absl/app.py", line 251 in _run_main
  File "/usr/local/lib/python3.6/dist-packages/absl/app.py", line 303 in run
  File "save_model.py", line 58 in <module>

Aborted (core dumped)

Any idea on this issue?