class_docstr: Implementation of Quantization-Aware Training (QAT) for PyTorch models.
  
  :param model: (Quantizable) Model that needs to be quantized
  :type model: torch.nn.Module
  :param train_loader: DataLoader used for training
  :type train_loader: torch.utils.data.DataLoader
  :param test_loader: DataLoader used for testing
  :type test_loader: torch.utils.data.DataLoader
  :param optimizer: Optimizer for training
  :type optimizer: torch.optim.*
  :param qconfig: Configuration used for quantization
  :type qconfig: Qconfig
  :param criterion: Loss function used for training
  :type criterion: Loss_fn
  :param device: Device used for training ("cpu" or "cuda")
  :type device: torch.device
docstr: Function used for training the model
  
  :param model (torch.nn.Module): Model that needs to be trained
  :param optimizer (torch.optim.*): Optimizer for training
  :param num_batches (int): Number of batches used for calibration
docstr_tok: ["Function","used","for","training","the","model",":","param","model","(","torch.nn.Module",")",":","Model","that","needs","to","be","trained",":","param","optimizer","(","torch.optim",".","*",")",":","Optimizer","for","training",":","param","num_batches","(","int",")",":","Number","of","batches","used","for","calibration"]
code: def _train_model(self, model, optimizer, num_batches):
      """
      Function used for training the model
  
      :param model (torch.nn.Module): Model that needs to be trained
      :param optimizer (torch.optim.*): Optimizer for training
      :param num_batches (int): Number of batches used for calibration
      """
  
      model.to(self.device)
      model.train()
  
      correct = 0
      epoch_loss = 0
      cnt = 0
      len_dataset = min(
          num_batches * self.train_loader.batch_size,
          len(self.train_loader.dataset),
      )
  
      for image, target in self.train_loader:
          image, target = image.to(self.device), target.to(self.device)
          output = model(image)
  
          if isinstance(output, tuple):
              output = output[0]
  
          loss = self.criterion(output, target)
  
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
          cnt += 1
          pred = output.argmax(dim=1, keepdim=True)
          correct += pred.eq(target.view_as(pred)).sum().item()
  
          epoch_loss += loss.item()
  
          if cnt >= num_batches:
              return epoch_loss, (correct / len_dataset)
  
      return epoch_loss, (correct / len_dataset)
  
file_path: KD_Lib/KD_Lib/Quantization/qat/qat.py
imports: from copy import deepcopy
  import torch
  from KD_Lib.Quantization.common import Quantizer
label_desc: Use image node_id 3 for calling the QAT_Quantizer obj's underlying member method code with example usage: obj._train_model(model, optimizer, num_batches) and returns: epoch_loss,  unknown, epoch_loss,  unknown