Update DNABERT_Local/bert_layers.py

DNABERT_Local/bert_layers.py

@@ -389,18 +389,23 @@ class BertEncoder(nn.Module):
             slopes_b = slopes_b[0::2][:n_heads - closest_power_of_2]
             return slopes_a + slopes_b
 
-        context_position = torch.arange(size, device=device)[:, None]
-        memory_position = torch.arange(size, device=device)[None, :]
+        # 🚨 THE ZEROGPU CPU-LOCK PATCH 🚨
+        # We force all ALiBi math to strictly use the CPU, completely bypassing the "meta" device collisions.
+        # The forward() function will automatically move this to the A100 GPU later when needed!
+        context_position = torch.arange(size, device='cpu')[:, None]
+        memory_position = torch.arange(size, device='cpu')[None, :]
         relative_position = torch.abs(memory_position - context_position)
+        
         # [n_heads, max_token_length, max_token_length]
-        relative_position = relative_position.unsqueeze(0).expand(
-            n_heads, -1, -1)
-        slopes = torch.Tensor(_get_alibi_head_slopes(n_heads)).to(device)
+        relative_position = relative_position.unsqueeze(0).expand(n_heads, -1, -1)
+        
+        slopes = torch.tensor(_get_alibi_head_slopes(n_heads), dtype=torch.float32, device='cpu')
         alibi = slopes.unsqueeze(1).unsqueeze(1) * -relative_position
+        
         # [1, n_heads, max_token_length, max_token_length]
         alibi = alibi.unsqueeze(0)
         assert alibi.shape == torch.Size([1, n_heads, size, size])
-
+        
         self._current_alibi_size = size
         self.alibi = alibi