diff --git a/examples/game_gomoku/code/board_agent.py b/examples/game_gomoku/code/board_agent.py
index d84051277..6cbef4ced 100644
--- a/examples/game_gomoku/code/board_agent.py
+++ b/examples/game_gomoku/code/board_agent.py
@@ -25,7 +25,7 @@
 EMPTY_PIECE = "0"
 
 
-def board2img(board: np.ndarray, output_dir: str) -> str:
+def board2img(board: np.ndarray, save_path: str) -> str:
     """Convert the board to an image and save it to the specified path."""
 
     size = board.shape[0]
@@ -63,9 +63,9 @@ def board2img(board: np.ndarray, output_dir: str) -> str:
     ax.set_xticklabels(range(size))
     ax.set_yticklabels(range(size))
     ax.invert_yaxis()
-    plt.savefig(output_dir, bbox_inches="tight", pad_inches=0.1)
+    plt.savefig(save_path, bbox_inches="tight", pad_inches=0.1)
     plt.close(fig)  # Close the figure to free memory
-    return output_dir
+    return save_path
 
 
 class BoardAgent(AgentBase):