Update modeling_chessbot.py

modeling_chessbot.py

@@ -653,33 +653,33 @@ class ChessBotModel(ChessBotPreTrainedModel):
         
         return selected_move
         
- def get_position_value(self, fen, device="cuda"):
-    """
-    Get the value evaluation for a given FEN position.
-    Returns the value vector [black_win_prob, draw_prob, white_win_prob]
-    """
-    x = torch.from_numpy(fen_to_tensor(fen)).to(device).to(torch.float32)
-    x = x.view(1, 1, 8, 8, 19)
-    
-    # Forward pass through the model to get value
-    with torch.no_grad():
-        # We need to run through the model layers to get to value_head
-        b, seq_len, _, _, emb = x.size()
-        x_processed = x.view(b * seq_len, 64, emb)
-        x_processed = self.linear1(x_processed)
-        x_processed = F.gelu(x_processed)
-        x_processed = self.layernorm1(x_processed)
-        x_processed = self.ma_gating(x_processed)
-        
-        pos_enc = self.positional(x_processed)
-        for i in range(self.num_layers):
-            x_processed = self.layers[i](x_processed, pos_enc)
+    def get_position_value(self, fen, device="cuda"):
+        """
+        Get the value evaluation for a given FEN position.
+        Returns the value vector [black_win_prob, draw_prob, white_win_prob]
+        """
+        x = torch.from_numpy(fen_to_tensor(fen)).to(device).to(torch.float32)
+        x = x.view(1, 1, 8, 8, 19)
         
-        value_logits = self.value_head_q(x_processed)
-        value_logits = value_logits.view(b, seq_len, 3)
-        value_logits = torch.softmax(value_logits, dim=-1)
-         
-    return value_logits.squeeze()  # Remove batch and sequence dimensions
+        # Forward pass through the model to get value
+        with torch.no_grad():
+            # We need to run through the model layers to get to value_head
+            b, seq_len, _, _, emb = x.size()
+            x_processed = x.view(b * seq_len, 64, emb)
+            x_processed = self.linear1(x_processed)
+            x_processed = F.gelu(x_processed)
+            x_processed = self.layernorm1(x_processed)
+            x_processed = self.ma_gating(x_processed)
+            
+            pos_enc = self.positional(x_processed)
+            for i in range(self.num_layers):
+                x_processed = self.layers[i](x_processed, pos_enc)
+            
+            value_logits = self.value_head_q(x_processed)
+            value_logits = value_logits.view(b, seq_len, 3)
+            value_logits = torch.softmax(value_logits, dim=-1)
+             
+        return value_logits.squeeze()  # Remove batch and sequence dimensions
 
     def get_batch_position_values(self, fens, device="cuda"):
         """