Update comment on shape in SymbolModality

PiperOrigin-RevId: 191759697
tensorflow · Apr 5, 2018 · b39d152 · b39d152
1 parent fc9335c
commit b39d152
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tensor2tensor/layers/modalities.py b/tensor2tensor/layers/modalities.py
@@ -95,11 +95,12 @@ def _get_weights(self, hidden_dim=None):
 
   def bottom_simple(self, x, name, reuse):
     with tf.variable_scope(name, reuse=reuse):
-      # Squeeze out the channels dimension.
+      # Ensure the inputs are 3-D
       if len(x.get_shape()) == 4:
         x = tf.squeeze(x, axis=3)
       while len(x.get_shape()) < 3:
         x = tf.expand_dims(x, axis=-1)
+
       var = self._get_weights()
       x = common_layers.dropout_no_scaling(
           x, 1.0 - self._model_hparams.symbol_dropout)