 public struct Dropout<Scalar>: Layer
     where Scalar: BinaryFloatingPoint & Differentiable & TensorFlowScalar,
           Scalar.RawSignificand : FixedWidthInteger {
-    @noDerivative public let dropProb: Double
-    @noDerivative public let mode: ModeRef
+    @noDerivative public let probability: Double
+    @noDerivative public let learningPhaseIndicator: LearningPhaseIndicator
+    // Workaround for TF-8
     var _unused: Tensor<Scalar>
 
-    public init(dropProbability: Double, modeRef: ModeRef) {
-        self.dropProb = dropProbability
-        self.mode = modeRef
-        self._unused = Tensor<Scalar>(Scalar(0))
-    }
-
-    @differentiable(wrt: (self, input))
-    public func applyTraining(to input: Tensor<Scalar>) -> Tensor<Scalar> {
-        return input.droppedOut(withProbability: dropProb)
-    }
-
-    @differentiable(wrt: (self, input))
-    public func applyInference(to input: Tensor<Scalar>) -> Tensor<Scalar> {
+    public init(
+        probability: Double,
+        learningPhaseIndicator: LearningPhaseIndicator
+    ) {
+        self.probability = probability
+        self.learningPhaseIndicator = learningPhaseIndicator
+        // Workaround for TF-8
+        self._unused = Tensor<Scalar>(0)
+    }
+
+    @differentiable(wrt: (self, input))
+    func applyTraining(to input: Tensor<Scalar>) -> Tensor<Scalar> {
+        return input.droppedOut(withProbability: probability)
+    }
+
+    @differentiable(wrt: (self, input))
+    func applyInference(to input: Tensor<Scalar>) -> Tensor<Scalar> {
         return input
     }
 
     @differentiable(wrt: (self, input), vjp: _vjpApplied(to:))
     public func applied(to input: Tensor<Scalar>) -> Tensor<Scalar> {
-        if mode.training {
+        if learningPhaseIndicator.training {
             return applyTraining(to: input)
         } else {
             return applyInference(to: input)

    public func _vjpApplied(to input: Tensor<Scalar>) ->
        (Tensor<Scalar>, (Tensor<Scalar>) ->
            (Dropout<Scalar>.CotangentVector, Tensor<Scalar>)) {
        if mode.training {
            return Swift.valueWithPullback(at: self, input) {
                $0.applyTraining(to: $1)
            }
        } else {
            return Swift.valueWithPullback(at: self, input) {
                $0.applyInference(to: $1)
            }
        }

        self.epsilon = epsilon
        self.runningMean = Parameter(Tensor(0))
        self.runningVariance = Parameter(Tensor(1))
        self.learningPhaseIndicator = learningPhaseIndicator
    }
}

