在gpu上在TensorFlow中服用FFT2D时的错误

Question

下面的代码在CPU上运行没有任何问题，但是当我更改为Colab中的GPU时，它无法计算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])

但是在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

我降级到Tensorflow 2.8.2，问题已经消失。