class_docstr: Wrapper class for Hugging Face Transformers PyTorch T5 model.
class_docstr_tok: ["Wrapper","class","for","Hugging","Face","Transformers","PyTorch","T5","model","."]
docstr: Train the model on the given Mixture or Task.
  
  Args:
    mixture_or_task_name: str, the name of the Mixture or Task to train on.
      Must be pre-registered in the global `t5.data.TaskRegistry` or
      `t5.data.MixtureRegistry.`
    steps: int, the total number of steps to train for.
    save_steps: int, the number of steps between checkpoint saves.
    sequence_length: dict of int, a dict mapping feature name to length.
    split: str or `tensorflow_datasets.Split`, the data split to load.
    batch_size: int, the number of padded sequences in each batch.
    optimizer: function that takes the model parameters as its sole argument.
      For example, to use an AdamW optimizer with a learning rate of 1e-4,
      you could pass in `functools.partial(transformers.AdamW, lr=1e-4)`.
    learning_rate_scheduler: optional function that takes in an optimizer as
      its sole argument. For example, to use a schedule that warms up the
      optimizer's learning rate after 100 steps, you could pass in
      `functools.partial(transformers.get_constant_schedule_with_warmup,
     num_warmup_steps=100)`.
docstr_tok: ["Train","the","model","on","the","given","Mixture","or","Task",".","Args",":","mixture_or_task_name",":","str",",","the","name","of","the","Mixture","or","Task","to","train","on",".","Must","be","pre-registered","in","the","global","`","t5.data.TaskRegistry","`","or","`","t5.data.MixtureRegistry.","`","steps",":","int",",","the","total","number","of","steps","to","train","for",".","save_steps",":","int",",","the","number","of","steps","between","checkpoint","saves",".","sequence_length",":","dict","of","int",",","a","dict","mapping","feature","name","to","length",".","split",":","str","or","`","tensorflow_datasets.Split","`",",","the","data","split","to","load",".","batch_size",":","int",",","the","number","of","padded","sequences","in","each","batch",".","optimizer",":","function","that","takes","the","model","parameters","as","its","sole","argument",".","For","example",",","to","use","an","AdamW","optimizer","with","a","learning","rate","of","1e-4",",","you","could","pass","in","`","functools.partial","(","transformers.AdamW",",","lr=1e-4",")","`",".","learning_rate_scheduler",":","optional","function","that","takes","in","an","optimizer","as","its","sole","argument",".","For","example",",","to","use","a","schedule","that","warms","up","the","optimizer","'s","learning","rate","after","100","steps",",","you","could","pass","in","`","functools.partial","(","transformers.get_constant_schedule_with_warmup",",","num_warmup_steps=100",")","`","."]
code: def train(
      self,
      mixture_or_task_name,
      steps,
      save_steps,
      sequence_length,
      split,
      batch_size,
      optimizer,
      learning_rate_scheduler=None,
  ):
      """Train the model on the given Mixture or Task.
  
      Args:
        mixture_or_task_name: str, the name of the Mixture or Task to train on.
          Must be pre-registered in the global `t5.data.TaskRegistry` or
          `t5.data.MixtureRegistry.`
        steps: int, the total number of steps to train for.
        save_steps: int, the number of steps between checkpoint saves.
        sequence_length: dict of int, a dict mapping feature name to length.
        split: str or `tensorflow_datasets.Split`, the data split to load.
        batch_size: int, the number of padded sequences in each batch.
        optimizer: function that takes the model parameters as its sole argument.
          For example, to use an AdamW optimizer with a learning rate of 1e-4,
          you could pass in `functools.partial(transformers.AdamW, lr=1e-4)`.
        learning_rate_scheduler: optional function that takes in an optimizer as
          its sole argument. For example, to use a schedule that warms up the
          optimizer's learning rate after 100 steps, you could pass in
          `functools.partial(transformers.get_constant_schedule_with_warmup,
         num_warmup_steps=100)`.
      """
      self._model.train()
      ds = _get_dataset(mixture_or_task_name, sequence_length, split)
      task = seqio.get_mixture_or_task(mixture_or_task_name)
      ds = tokens_to_batches(
          ds,
          sequence_length,
          batch_size,
          tuple(task.output_features),
          task,
      )
      # Repeat dataset forever
      ds = itertools.cycle(ds)
      optimizer = optimizer(self._model.parameters())
      if learning_rate_scheduler:
          learning_rate_scheduler = learning_rate_scheduler(optimizer)
  
      now = time.time()
      for train_step, batch in enumerate(itertools.islice(ds, steps)):
          if not train_step % save_steps:
              # TODO(craffel): Consider saving optimizer and scheduler state.
              logging.info("Saving checkpoint for step %s", self._step)
              self.save_checkpoint(self._step)
  
          self._model.zero_grad()
          outputs = self._model(
              input_ids=self.to_tensor(batch["inputs"]),
              attention_mask=self.to_tensor(batch["inputs_mask"]),
              decoder_attention_mask=self.to_tensor(
                  batch["targets_mask"]
              ),
              labels=self.to_tensor(batch["targets"]),
          )
          loss = outputs[0]
          loss.backward()
          optimizer.step()
          if learning_rate_scheduler:
              learning_rate_scheduler.step()
  
          self._writer.add_scalar(
              "loss", loss.detach().cpu().numpy(), self._step
          )
          self._writer.add_scalar(
              "step/s", 1 / (time.time() - now), self._step
          )
          now = time.time()
          self._step += 1
  
      logging.info("Saving final checkpoint for step %s", self._step)
      self.save_checkpoint(self._step)
  
imports: import functools
  import itertools
  import os
  import re
  import time
  from absl import logging
  import mesh_tensorflow.transformer.dataset
  import seqio
  import t5.data
  from t5.models import utils
  from t5.models.t5_model import T5Model
  import tensorflow.compat.v1
  import tensorflow_datasets
  import torch
  import torch.utils.tensorboard
label_desc: Use image node_id 9 for calling the HfPyTorchModel obj's underlying member method code with example usage: obj.train(mixture_or_task_name, steps, save_steps, sequence_length, split, batch_size, optimizer, learning_rate_scheduler) without return types