Update utils/engine.py

utils/engine.py

@@ -8,6 +8,7 @@ from dataclasses import dataclass, field
 from functools import partial
 import time
 import os
+import spaces
 
 try:
     from .mapping import UCI_MOVE_TO_IDX, IDX_TO_UCI_MOVE
@@ -162,7 +163,7 @@ class Engine:
                 raise ValueError(f"Invalid engine path or engine not found: {e}")
         else:
             raise ValueError("Invalid engine type. Choose 'chessformer' or 'stockfish'.")
-
+    
     def get_invalid_mask(self, boards: List[chess.Board]) -> torch.Tensor:
         bs = len(boards)
         possible_moves = len(UCI_MOVE_TO_IDX)
@@ -617,13 +618,16 @@ class Engine:
             
             return results
 
+    @spaces.GPU
     def move(self, board: chess.Board, return_perplexity: bool=False) -> Tuple[str, float]:
         if self.type == "stockfish":
             return self._stockfish_move(board, return_perplexity)
         elif self.type == "chessformer":
+            self.model.to("cuda")
             return self._chessformer_move(board, return_perplexity)
         else:
             raise ValueError(f"Invalid engine type: {self.type}")
+        self.model.to("cpu")
         
     def batch_move(self, boards: List[chess.Board]) -> List[Tuple[str, float]]:
         if self.type == "stockfish":
@@ -632,7 +636,8 @@ class Engine:
             return self._batch_chessformer_move(boards)
         else:
             raise ValueError(f"Invalid engine type: {self.type}")
-        
+
+    @spaces.GPU
     def analyze_position(self, board: chess.Board) -> Optional[float]:
         """
         Analyzes the given **single board** position using the engine.
@@ -683,8 +688,11 @@ class Engine:
             fen = board.fen()
             count_tensor = self.compute_repetition([board.copy(stack=True)])
 
+            gpu_model = self.model.to("cuda")
             with torch.no_grad():
-                _, value = self.model([fen],count_tensor)
+                _, value = gpu_model([fen],count_tensor)
+
+            self.model = self.model.to("cpu")
 
             value = value.item()
             return value if board.turn == chess.WHITE else -value