[ModelZoo] Support variable type BF16 in DCN model. (DeepRec-AI#480)

15861329502 · Oct 13, 2022 · 28ac395 · 28ac395
1 parent cd0c16a
commit 28ac395
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/modelzoo/dcn/train.py b/modelzoo/dcn/train.py
@@ -166,10 +166,10 @@ def _cross_net(self, cross_input, layer_num=2, layer_name=''):
                                    reuse=tf.AUTO_REUSE) as cross_layer_scope:
                 w = tf.get_variable(
                     name=layer_name+'_w',
-                    dtype=tf.float32,
+                    dtype=cross_input.dtype,
                     shape=(last_dim),
                 )
-                b = tf.get_variable(name=layer_name+'_b', dtype=tf.float32, shape=(last_dim))
+                b = tf.get_variable(name=layer_name+'_b', dtype=cross_input.dtype, shape=(last_dim))
                 xw = tf.reduce_sum(x * w, axis=1, keepdims=True)
                 x = tf.math.add(tf.math.add(x0 * xw, b), x)