class_docstr: Single training iteration loop of the model.
  
  Performs: data read, forward, loss computation, backward, optimizer step, parameter updates.
  
  Various intermediate steps are also performed:
  - logging the training loss, training eta, LR, etc to loggers
  - logging to tensorboard,
  - performing any self-supervised method specific operations (like in MoCo approach, the
  momentum encoder is updated), computing the scores in swav
  - checkpointing model if user wants to checkpoint in the middle
  of an epoch
docstr: Given the input batch from the dataloader, verify the input is
  as expected: the input data and target data is present in the
  batch.
  In case of multi-input trainings like PIRL, make sure the data
  is in right format i.e. the multiple input should be nested
  under a common key "input".
docstr_tok: ["Given","the","input","batch","from","the","dataloader",",","verify","the","input","is","as","expected",":","the","input","data","and","target","data","is","present","in","the","batch",".","In","case","of","multi-input","trainings","like","PIRL",",","make","sure","the","data","is","in","right","format","i.e",".","the","multiple","input","should","be","nested","under","a","common","key","``","input","''","."]
code: def construct_sample_for_model(task: ClassyTask):
      """
      Given the input batch from the dataloader, verify the input is
      as expected: the input data and target data is present in the
      batch.
      In case of multi-input trainings like PIRL, make sure the data
      is in right format i.e. the multiple input should be nested
      under a common key "input".
      """
      # Process next sample
      with PerfTimer("read_sample", task.perf_stats):
          batch_data = next(task.data_iterator)
  
      sample_key_names = task.data_and_label_keys
      inp_key, target_key = (
          sample_key_names["input"],
          sample_key_names["target"],
      )
      all_keys = inp_key + target_key
  
      assert len(inp_key) + len(target_key) <= len(
          batch_data
      ), "Number of input and target keys in batch and train config don't match."
  
      # every input should be a list. The list corresponds to various data sources
      # and hence could be used to handle several data modalities.
      for key in all_keys:
          assert isinstance(
              batch_data[key], list
          ), f"key: {key} input is not a list"
          assert (
              len(batch_data[key]) == 1
          ), f"Please modify your train step to handle multi-modal input: key {key}"
  
      # single input case
      if (
          len(sample_key_names["input"]) == 1
          and len(sample_key_names["target"]) == 1
      ):
          sample = {
              "input": batch_data[inp_key[0]][0],
              "target": batch_data[target_key[0]][0],
              "data_valid": batch_data["data_valid"][0],
          }
  
      # multi-input case (example in PIRL, we pass image and patches both).
      # we nest all these under the sample["input"]
      elif len(sample_key_names["input"]) > 1:
          sample = {"input": {}, "target": {}, "data_valid": None}
          for key in inp_key:
              sample["input"][key] = batch_data[key][0]
  
          if len(target_key) > 1:
              for key in target_key:
                  sample["target"][key] = batch_data[key][0]
          else:
              sample["target"] = batch_data[target_key[0]][0]
          sample["data_valid"] = batch_data["data_valid"][0]
  
      # Copy the other keys as-is, method dependent
      # - But avoid to erase the standard keys
      for k in batch_data.keys():
          if k not in all_keys and k not in sample:
              sample[k] = batch_data[k]
  
      return sample
  
imports: import contextlib
  from types import SimpleNamespace
  from typing import Any, Dict
  import torch
  from classy_vision.generic.distributed_util import all_reduce_mean
  from classy_vision.tasks import ClassyTask
  from classy_vision.tasks.classification_task import AmpType
  from vissl.hooks import SSLClassyHookFunctions
  from vissl.trainer.train_steps import register_train_step
  from vissl.utils.activation_checkpointing import manual_gradient_all_reduce, manual_sync_params
  from vissl.utils.fsdp_utils import is_fsdp_model
  from vissl.utils.misc import is_apex_available, torch_version
  from vissl.utils.perf_stats import PerfTimer
  from vissl.utils.profiler import record_function
label_desc: Use image node_id 2 for calling the StandardTrainStep obj's underlying member method code with example usage: obj.construct_sample_for_model(task) and returns: sample