Update model/modules.py

model/modules.py

@@ -208,13 +208,19 @@ def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_resca
 
 
 def get_pos_embed_indices(start, length, max_pos, scale=1.0):
-    # length = length if isinstance(length, int) else length.max()
-    scale = scale * torch.ones_like(start, dtype=torch.float32)  # in case scale is a scalar
+    """
+    Generate positional embedding indices with bfloat16 for computations.
+    """
+    # Convert scale to a tensor with bfloat16
+    scale = scale * torch.ones_like(start, dtype=torch.bfloat16)
+    
+    # Compute positions using bfloat16
     pos = (
         start.unsqueeze(1)
-        + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()
+        + (torch.arange(length, device=start.device, dtype=torch.bfloat16).unsqueeze(0) * scale.unsqueeze(1)).long()
     )
-    # avoid extra long error.
+    
+    # Ensure positions do not exceed max_pos
     pos = torch.where(pos < max_pos, pos, max_pos - 1)
     return pos