Quick fix for 0.1.7 keras-core release (#1251)

keras-team · Sep 21, 2023 · 0b0b9d5 · 0b0b9d5
1 parent 5cf2c5e
commit 0b0b9d5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/keras_nlp/models/xlnet/xlnet_content_and_query_embedding.py b/keras_nlp/models/xlnet/xlnet_content_and_query_embedding.py
@@ -67,12 +67,12 @@ def positional_embedding(self, pos_seq, inv_freq, bsz=None):
 
     def relative_positional_encoding(self, qlen, klen, bsz=None, clamp_len=-1):
         """create relative positional encoding."""
-        freq_seq = ops.arange(0, self.hidden_dim, 2.0)
+        freq_seq = ops.arange(0, self.hidden_dim, 2.0, dtype=self.compute_dtype)
         inv_freq = 1 / (10000 ** (freq_seq / self.hidden_dim))
 
         beg, end = klen, -qlen
 
-        fwd_pos_seq = ops.arange(beg, end, -1.0)
+        fwd_pos_seq = ops.arange(beg, end, -1.0, dtype=self.compute_dtype)
         if clamp_len > 0:
             fwd_pos_seq = ops.clip(
                 fwd_pos_seq, x_min=-clamp_len, x_max=clamp_len