Update model/modules.py

model/modules.py

@@ -212,12 +212,12 @@ def get_pos_embed_indices(start, length, max_pos, scale=1.0):
     Generate positional embedding indices with bfloat16 for computations.
     """
     # Convert scale to a tensor with bfloat16
-    scale = scale * torch.ones_like(start, dtype=torch.bfloat16)
+    scale = scale * torch.ones_like(start, dtype=torch.float32)
     
     # Compute positions using bfloat16
     pos = (
         start.unsqueeze(1)
-        + (torch.arange(length, device=start.device, dtype=torch.bfloat16).unsqueeze(0) * scale.unsqueeze(1)).long()
+        + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()
     )
     
     # Ensure positions do not exceed max_pos