Error when taking fft2d in TensorFlow on GPU

Question

The code below runs on CPU without any problems, But when I change to GPU in colab it fails to calculate the fft2d.

import tensorflow as tf

sample_fft_input = tf.random.uniform((2, 10, 20))
sfi = tf.cast(sample_fft_input , tf.complex64)
sfi = tf.math.real(tf.signal.fft2d(sfi))
print(sfi.shape) -> TensorShape([2, 10, 20])

But on GPU:

---------------------------------------------------------------------------
InternalError                             Traceback (most recent call last)
<ipython-input-41-094e0f7d5037> in <module>()
      3 sample_fft_input = tf.random.uniform((2, 10, 20))
      4 sfi = tf.cast(sample_fft_input, tf.complex64)
----> 5 sfi = tf.math.real(tf.signal.fft2d(sfi))
      6 print(sfi.shape)

1 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py in raise_from_not_ok_status(e, name)
   7162 def raise_from_not_ok_status(e, name):
   7163   e.message += (" name: " + name if name is not None else "")
-> 7164   raise core._status_to_exception(e) from None  # pylint: disable=protected-access
   7165 
   7166 

InternalError: fft failed : type=1 in.shape=[2,10,20] [Op:FFT2D]

Answer 1

I downgraded to tensorflow 2.8.2 and the problem is gone.