docstr: Train model with a batch of tasks, update meta-model once.
  
  Args:
      data_batch: a batch of tasks, each task contains support set and query set
      current_epoch
  
  Returns:
      Average MSE loss on query set
docstr_tok: ["Train","model","with","a","batch","of","tasks",",","update","meta-model","once",".","Args",":","data_batch",":","a","batch","of","tasks",",","each","task","contains","support","set","and","query","set","current_epoch","Returns",":","Average","MSE","loss","on","query","set"]
code: def train_one_iter(self, data_batch, current_epoch):
      """Train model with a batch of tasks, update meta-model once.
  
      Args:
          data_batch: a batch of tasks, each task contains support set and query set
          current_epoch
  
      Returns:
          Average MSE loss on query set
      """
      (
          x_support_set,
          y_support_set,
          x_query_set,
          y_query_set,
      ) = data_batch
  
      x_support_set = x_support_set.to(self.device)
      y_support_set = y_support_set.to(self.device)
      x_query_set = x_query_set.to(self.device)
      y_query_set = y_query_set.to(self.device)
  
      self.model.zero_grad()
      if self.learning_rate_scheduler:
          self.scheduler.step(current_epoch)
      batch_losses = []
      for x_support, y_support, x_query, y_query in zip(
          x_support_set, y_support_set, x_query_set, y_query_set
      ):
          task_loss = 0
          per_step_loss_importance_vectors = (
              self._get_per_step_loss_importance_vector(current_epoch)
          )
          model_weights = self.model.get_weights()
  
          for num_step in range(self.num_updates_per_iter):
              support_loss = self._caculate_loss(
                  x_support, y_support, model_weights
              )
  
              model_weights = self._inner_loop_update(
                  loss=support_loss,
                  weights=model_weights,
                  use_second_order=self.second_order,
                  current_update_step=num_step,
              )
  
              if (
                  self.use_multi_step_loss_optimization
                  and current_epoch < self.multi_step_loss_num_epochs
              ):
                  query_loss = self._caculate_loss(
                      x=x_query, y=y_query, weights=model_weights
                  )
                  task_loss = (
                      task_loss
                      + per_step_loss_importance_vectors[num_step]
                      * query_loss
                  )
  
              elif num_step == (self.num_updates_per_iter - 1):
                  query_loss = self._caculate_loss(
                      x=x_query, y=y_query, weights=model_weights
                  )
                  task_loss = query_loss
  
          batch_losses.append(task_loss)
  
      avg_batch_loss = sum(batch_losses) / len(batch_losses)
  
      self._outer_loop_update(avg_batch_loss)
  
      return avg_batch_loss.cpu().item()
  
file_path: PARL/benchmark/torch/maml++/maml_algorithm.py
imports: import torch
  from torch import nn
  from torch import optim
  import torch.nn.functional
  import numpy
  import parl
label_desc: Use image node_id 2 for calling the MAML obj's underlying member method code with example usage: obj.train_one_iter(data_batch, current_epoch) and returns: avg_batch_loss