class_docstr: Trainer for the AlphaTensor model. The trainer does not require an
  explicit loss since the loss is computed by the model itself. The trainer
  is responsible for both the training step and the acting one, storing
  acting performance in a buffer.
class_docstr_tok: ["Trainer","for","the","AlphaTensor","model",".","The","trainer","does","not","require","an","explicit","loss","since","the","loss","is","computed","by","the","model","itself",".","The","trainer","is","responsible","for","both","the","training","step","and","the","acting","one",",","storing","acting","performance","in","a","buffer","."]
docstr: Executes a single training step by optimizing the current model
  parameters.
docstr_tok: ["Executes","a","single","training","step","by","optimizing","the","current","model","parameters","."]
code: def train_step(self):
      """Executes a single training step by optimizing the current model
      parameters."""
      self.dataset.recompute_synthetic_indexes()
      self.model.train()
      total_loss = 0
      dl = DataLoader(
          self.dataset, batch_size=self.batch_size, shuffle=True
      )
      print("Training AlphaTensor")
      for states, scalars, policies, rewards in tqdm.tqdm(dl):
          loss_policy, loss_value = self.model(
              states, scalars, policies, rewards
          )
          loss = self.alpha * loss_policy + self.beta * loss_value
          self.optimizer.zero_grad()
          loss.backward()
          self.optimizer.step()
          total_loss += loss.item()
      print(f"Total loss: {total_loss}")
  
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 the Trainer obj's underlying member method code with example usage: obj.train_step() without return types