diff --git a/‎tests/test_keras.py‎
Lines changed: 18 additions & 26 deletions b/‎tests/test_keras.py‎
Lines changed: 18 additions & 26 deletions
@@ -5,7 +5,7 @@
 import pandas as pd
 
 from keras.models import Sequential
-from keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPooling2D, CuDNNLSTM
+from keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPooling2D, LSTM
 from keras.optimizers import RMSprop, SGD
 from keras.utils.np_utils import to_categorical
 
@@ -62,29 +62,21 @@ def test_conv2d(self):
 
         model.evaluate(x_test, y_test, batch_size=32)
 
-    # Tensorflow 2.0 doesn't support the contrib package.
-    #
-    # Error:
-    #   from tensorflow.contrib.cudnn_rnn.python.ops import cudnn_rnn_ops
-    # ModuleNotFoundError: No module named 'tensorflow.contrib'
-    #
-    # tf.keras should be used instead until this is fixed.
-    # @gpu_test
-    # def test_cudnn_lstm(self):
-    #     x_train = np.random.random((100, 100, 100))
-    #     y_train = keras.utils.to_categorical(np.random.randint(10, size=(100, 1)), num_classes=10)
-    #     x_test = np.random.random((20, 100, 100))
-    #     y_test = keras.utils.to_categorical(np.random.randint(10, size=(20, 1)), num_classes=10)
-
-    #     sgd = SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)
-
-    #     model = Sequential()
-    #     model.add(CuDNNLSTM(32, return_sequences=True, input_shape=(100, 100)))
-    #     model.add(Flatten())
-    #     model.add(Dense(10, activation='softmax'))
-
-
-    #     model.compile(loss='categorical_crossentropy', optimizer=sgd)
-    #     model.fit(x_train, y_train, batch_size=32, epochs=1)
-    #     model.evaluate(x_test, y_test, batch_size=32)
+    def test_lstm(self):
+        x_train = np.random.random((100, 100, 100))
+        y_train = keras.utils.to_categorical(np.random.randint(10, size=(100, 1)), num_classes=10)
+        x_test = np.random.random((20, 100, 100))
+        y_test = keras.utils.to_categorical(np.random.randint(10, size=(20, 1)), num_classes=10)
+
+        sgd = SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)
+
+        model = Sequential()
+        model.add(LSTM(32, return_sequences=True, input_shape=(100, 100)))
+        model.add(Flatten())
+        model.add(Dense(10, activation='softmax'))
+
+
+        model.compile(loss='categorical_crossentropy', optimizer=sgd)
+        model.fit(x_train, y_train, batch_size=32, epochs=1)
+        model.evaluate(x_test, y_test, batch_size=32)