ValueError: exactly one of keep_prob and rate should be set #912

[ayush714]

When I am fine tuning, I am getting below error:-

~/.local/lib/python3.8/site-packages/mesh_tensorflow/transformer/transformer.py in call_simple(self, inputs, targets, compute_loss, mode, variable_dtype, sequence_id, subsequence_id, position, encoder_output, encoder_sequence_id, encoder_inputs, shared_params, layer_outputs, encoder_layer_outputs, num_microbatches)
   1035         num_microbatches=num_microbatches)
   1036     with tf.variable_scope(self.name):
-> 1037       logits = self._call_internal(context, inputs, targets)
   1038     if compute_loss:
   1039       loss = mtf.add_n(context.losses)

~/.local/lib/python3.8/site-packages/mesh_tensorflow/transformer/transformer.py in _call_internal(self, context, inputs, targets)
    854           name="embedding",
    855           ensemble_dim=self.ensemble_dim)
--> 856     inputs = mtf.dropout(inputs, context.train, rate=self.token_dropout_rate)
    857     x = vocab_embedding.ids_to_embedding(inputs, context)
    858     if self.positional_embedding or self.sinusoid_positional_embedding:

~/.local/lib/python3.8/site-packages/mesh_tensorflow/ops.py in dropout(x, keep_prob, rate, noise_shape, name)
   5820   Args:
   5821     mesh: a Mesh
-> 5822     shape: a Shape
   5823     **kwargs: keyword args for tf.random.normal, except seed
   5824 

**ValueError: exactly one of keep_prob and rate should be set**

[ayush714]

@adarob @craffel Please take a look at this issue. Still I am not able to debug