Deeplab Loss is inf or nan RTX3090

1. Describe the bug

When I try to train a Deeplab model, I get Invalid argument: Loss is inf or nan. : Tensor had NaN values [[{{node CheckNumerics}}]]. Training does not start.

This happens when I try it on a workstation with RTX 3090. I have tested the SAME environment on a GTX 1060 Windows and on a GTX 1070 ubuntu 18.04 and it works without any problem.

2. Expected behavior

Training should run as normal

3. System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04
• Mobile device name if the issue happens on a mobile device:
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 1.15 GPU
• Python version: 3.7
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version: CUDA 10.0 (conda cudatoolkit)
• GPU model and memory: GTX 3090