Fix StochasticDepth layer error in training mixed_float16 (#2450)

* Update stochastic_depth_test.py * Update stochastic_depth.py * Update stochastic_depth.py * Update stochastic_depth.py
tensorflow · Jul 27, 2021 · e9846e8 · e9846e8
1 parent 3e87219
commit e9846e8
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/tensorflow_addons/layers/stochastic_depth.py b/tensorflow_addons/layers/stochastic_depth.py
@@ -65,7 +65,9 @@ def call(self, x, training=None):
         shortcut, residual = x
 
         # Random bernoulli variable indicating whether the branch should be kept or not or not
-        b_l = tf.keras.backend.random_bernoulli([], p=self.survival_probability)
+        b_l = tf.keras.backend.random_bernoulli(
+            [], p=self.survival_probability, dtype=self._compute_dtype_object
+        )
 
         def _call_train():
             return shortcut + b_l * residual

diff --git a/tensorflow_addons/layers/tests/stochastic_depth_test.py b/tensorflow_addons/layers/tests/stochastic_depth_test.py
@@ -47,7 +47,9 @@ def test_with_mixed_precision_policy():
     residual = np.asarray([[0.2, 0.4, 0.5]])
 
     output = StochasticDepth()([shortcut, residual])
+    assert output.dtype == policy.compute_dtype
 
+    output = StochasticDepth()([shortcut, residual], training=True)
     assert output.dtype == policy.compute_dtype