def _train_batch_(self):
    (data, target) = self.train_loader_it.next()
    (data, target) = (data.to(self.device), target.to(self.device))
    self.current_batch_size = data.size()[0]
    output = self.model(data)
    loss = self.loss_function(output, target)
    with backpack(BatchGrad()):
        loss.mean().backward()
    loss_value = loss.mean()
    reward = self._get_validation_loss()
    self.loss_batch = loss
    self.current_training_loss = torch.unsqueeze(loss_value.detach(), dim=0)
    self.train_batch_index += 1
    return reward
def _train_batch_(self):
    (data, target) = self.train_loader_it.next()
    (data, target) = (data.to(self.device), target.to(self.device))
    self.current_batch_size = data.size()[0]
    output = self.model(data)
    loss = self.loss_function(output, target)
    with backpack(BatchGrad()):
        loss.mean().backward()
    loss_value = loss.mean()
    reward = self._get_validation_loss()
    self.loss_batch = loss
    self.current_training_loss = torch.unsqueeze(loss_value.detach(), dim=0)
    self.train_batch_index += 1
    return reward