Remove dead code from shakespeare_main model

PiperOrigin-RevId: 272077584

official/staging/shakespeare/shakespeare_main.py

@@ -138,15 +138,8 @@ def build_model(vocab_size,
   Returns:
     A Keras Model.
   """
-  # In V1 there is a separate class for CuDNN. In V2 the LSTM class will use
-  # CuDNN automatically if applicable.
-  if use_cudnn and not keras_utils.is_v2_0():
-    LSTM = tf.compat.v1.CuDNNLSTM
-  else:
-    # The LSTM call was rewritten to be more efficient in 2.0. However because
-    # we want to compare the performance of the two runtimes, we force both
-    # V1 and V2 to use the more efficient implementation.
-    LSTM = functools.partial(tf.keras.layers.LSTM, implementation=2)
+  assert keras_utils.is_v2_0()
+  LSTM = functools.partial(tf.keras.layers.LSTM, implementation=2)
 
   # By indirecting the activation through a lambda layer, the logic to dispatch
   # to CuDNN in V2 doesn't trigger and we force the LSTM to run in non-CuDNN