if st == "element":
noise_shape = None
elif st == "item":
noise_shape = tf.concat([tf.shape(dropout_input)[:ids_rank],
tf.ones([self._dim_rank], tf.int32)],
axis=0)
elif st == "item_type":
After Change
if st == "element":
noise_shape = None
elif st == "item":
assert dropout_input is not None
assert ids_rank is not None
noise_shape = (dropout_input.get_shape().as_list()[:ids_rank]
+ [1] * self._dim_rank)
elif st == "item_type":