code: def swap_data(
      states: List[torch.Tensor],
      actions: List[torch.Tensor],
  ):
      """Swaps the last action with a random one and updates the states
      accordingly for a single game.
  
      Args:
          states (List[torch.Tensor]): All the states for a single game.
          actions (List[torch.Tensor]): All the actions through the game.
      """
      last_action = actions[-1]
      swap_index = torch.randint(0, len(states) - 1, (1,)).item()
      actions[-1] = actions[swap_index]
      actions[swap_index] = last_action
  
      actual_state = states[swap_index]
      for i in range(swap_index + 1, len(states) + 1):
          prev_action = actions[i - 1]
          triplet = map_action_to_triplet(
              prev_action, vector_size=actual_state.shape[-1]
          )
          vector_size = actual_state.shape[-1] // 3
          bs = actual_state.shape[0]
          u = triplet[:, :vector_size].reshape(bs, -1, 1, 1)
          v = triplet[
              :, vector_size : 2 * vector_size
          ].reshape(  # noqa E203
              bs, 1, -1, 1
          )
          w = triplet[:, 2 * vector_size :].reshape(
              bs, 1, 1, -1
          )  # noqa E203
          reduced_state = u * v * w
          fut_state = actual_state[:, 0] - reduced_state
          new_state = actual_state[:, 1:].roll(1, dims=1)
          new_state[:, 0] = reduced_state
          actual_state = torch.cat([fut_state, new_state], dim=1)
          states[i] = actual_state
      return states, actions
  
file_path: nebuly/optimization/open_alpha_tensor/open_alpha_tensor/core/training.py
imports: from pathlib import Path
  from typing import Tuple, List
  import torch.optim
  import tqdm
  from torch.utils.data import DataLoader
  from open_alpha_tensor.config import BASE_CHECKPOINT_DATA_DIR, BASE_CHECKPOINT_DIR
  from open_alpha_tensor.core.actors.stage import actor_prediction
  from open_alpha_tensor.core.data.basis_change import ChangeOfBasis
  from open_alpha_tensor.core.data.dataset import TensorGameDataset
  from open_alpha_tensor.core.data.generation import f_prob_distribution
  from open_alpha_tensor.core.data.utils import map_action_to_triplet
  from open_alpha_tensor.core.modules.alpha_tensor import AlphaTensorModel
label_desc: Use image node_id 2 for calling a global function with example usage: swap_data(states, actions) and returns: states,  actions