diff --git a/physics_planning_games/board_games/tic_tac_toe_logic.py b/physics_planning_games/board_games/tic_tac_toe_logic.py
index 3705cde..6e205a7 100644
--- a/physics_planning_games/board_games/tic_tac_toe_logic.py
+++ b/physics_planning_games/board_games/tic_tac_toe_logic.py
@@ -117,7 +117,7 @@ class TicTacToeRandomOpponent(logic_base.Opponent):
     valid_moves = game_logic.open_spiel_state.legal_actions()
     assert valid_moves
     move = random_state.choice(valid_moves)
-    row, col = np.unravel_index(move, dims=(3, 3))
+    row, col = np.unravel_index(move, shape=(3, 3))
     return SingleMarkerAction(row=row, col=col)
 
 
@@ -196,7 +196,7 @@ def numpy_array_to_open_spiel_state(board_state):
 
 
 def open_spiel_move_to_single_marker_action(action):
-  row, col = np.unravel_index(action, dims=(3, 3))
+  row, col = np.unravel_index(action, shape=(3, 3))
   return SingleMarkerAction(row=row, col=col)