Update utils/engine.py

optimize zero gpu acquirement

utils/engine.py

@@ -440,12 +440,21 @@ class Engine:
         else:
             return move_uci, final_value
 
+    @spaces.GPU
     def _chessformer_move(self, board: chess.Board, return_perplexity: bool=False, verbose: bool=False) -> Tuple[str,float]:
         """Get move from chessformer with optional search enhance"""
+        self.device = torch.device("cuda")
+        self.model.to(torch.device("cuda"))
         if self.depth == 0:
-            return self._raw_chessformer_move(board,return_perplexity)
+            result = self._raw_chessformer_move(board,return_perplexity)
+            self.model.to(torch.device("cpu"))
+            self.device = torch.device("cpu")
+            return result
         else:
-            return self._search_enhanced_move(board,return_perplexity,verbose)
+            result = self._search_enhanced_move(board,return_perplexity,verbose)
+            self.model.to(torch.device("cpu"))
+            self.device = torch.device("cpu")
+            return result
 
     def _stockfish_move(self, board: chess.Board, return_perplexity: bool=False) -> Tuple[str,float]:
         """Get best move from stockfish"""
@@ -618,18 +627,13 @@ 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.device = torch.device("cuda")
-            self.model.to(torch.device("cuda"))
             return self._chessformer_move(board, return_perplexity)
         else:
             raise ValueError(f"Invalid engine type: {self.type}")
-        self.model.to(torch.device("cpu"))
-        self.device = torch.device("cpu")
         
     def batch_move(self, boards: List[chess.Board]) -> List[Tuple[str, float]]:
         if self.type == "stockfish":
@@ -640,6 +644,26 @@ class Engine:
             raise ValueError(f"Invalid engine type: {self.type}")
 
     @spaces.GPU
+    def _analyze_position_gpu(self, board: chess.Board) -> Optional[float]:
+        """
+        Only acquire ZeroGPU when model forward pass is needed.
+        """
+        fen = board.fen()
+        
+        self.model = self.model.to(torch.device("cuda"))
+        self.device = torch.device("cuda")
+        
+        count_tensor = self.compute_repetition([board.copy(stack=True)])
+        
+        with torch.no_grad():
+            _, value = self.model([fen],count_tensor)
+
+        self.model = self.model.to("cpu")
+        self.device = torch.device("cpu")
+
+        value = value.item()
+        return value if board.turn == chess.WHITE else -value
+        
     def analyze_position(self, board: chess.Board) -> Optional[float]:
         """
         Analyzes the given **single board** position using the engine.
@@ -687,22 +711,7 @@ class Engine:
             
         
         elif self.type == "chessformer":
-            fen = board.fen()
-            
-            self.model = self.model.to(torch.device("cuda"))
-            self.device = torch.device("cuda")
-            
-            count_tensor = self.compute_repetition([board.copy(stack=True)])
-            
-            with torch.no_grad():
-                _, value = self.model([fen],count_tensor)
-
-            self.model = self.model.to("cpu")
-            self.device = torch.device("cpu")
-
-            value = value.item()
-            return value if board.turn == chess.WHITE else -value
-    
+            return self._analyze_position_gpu(board)
         else:
             raise ValueError(f"Invalid engine type.")