Merge pull request #16 from leehosu01/hotfix/mixed_float16

Update model.py, mixed_float16 feature
rishigami · Jan 23, 2022 · bd4d9f1 · bd4d9f1
2 parents 8986ca7 + 137ff4b
commit bd4d9f1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/swintransformer/model.py b/swintransformer/model.py
@@ -103,7 +103,7 @@ def call(self, x, mask=None):
         if mask is not None:
             nW = mask.get_shape()[0]  # tf.shape(mask)[0]
             attn = tf.reshape(attn, shape=[-1, nW, self.num_heads, N, N]) + tf.cast(
-                tf.expand_dims(tf.expand_dims(mask, axis=1), axis=0), tf.float32)
+                tf.expand_dims(tf.expand_dims(mask, axis=1), axis=0), attn.dtype)
             attn = tf.reshape(attn, shape=[-1, self.num_heads, N, N])
             attn = tf.nn.softmax(attn, axis=-1)
         else: