Is it possible to use CuDNNLSTM with Google Colab's TPU?

Viewed 928

I am able to do this with their GPU, but with their TPU it retrieves me an error... Does anybody around here know what I'm missing, please? Does it make sense to actually use the TPU with CuDNNLSTM? Or is CuDNNLSTM just tailored for GPU? Thanks a lot in advance.

1 Answers

keras.layers.CuDNNLSTM is only supported on GPUs. But in Tensorflow 2 built-in LSTM and GRU layers have been updated to leverage CuDNN kernels by default when a GPU is available.

Below is the details from Performance optimization and CuDNN kernels:

In TensorFlow 2.0, the built-in LSTM and GRU layers have been updated to leverage CuDNN kernels by default when a GPU is available. With this change, the prior keras.layers.CuDNNLSTM/CuDNNGRU layers have been deprecated, and you can build your model without worrying about the hardware it will run on.

You can just use the built-in LSTM layer: tf.keras.layers.LSTM and it will work on both TPUs and GPUs.

Related