diff --git a/model/modeling.py b/model/modeling.py
index f2a2030..54d4acd 100644
--- a/model/modeling.py
+++ b/model/modeling.py
@@ -364,8 +364,7 @@ def dropout(input_tensor, dropout_prob):
 
   Args:
     input_tensor: float Tensor.
-    dropout_prob: Python float. The probability of dropping out a value (NOT of
-      *keeping* a dimension as in `tf.nn.dropout`).
+    dropout_prob: Python float. The probability of dropping out a value.
 
   Returns:
     A version of `input_tensor` with dropout applied.
@@ -373,7 +372,7 @@ def dropout(input_tensor, dropout_prob):
   if dropout_prob is None or dropout_prob == 0.0:
     return input_tensor
 
-  output = tf.nn.dropout(input_tensor, 1.0 - dropout_prob)
+  output = tf.nn.dropout(input_tensor, rate=dropout_prob)
   return output