Run single batch of data

def feed_batch(self, batch_info, data, target):
        """ Run single batch of data """
        data, target = data.to(self.device), target.to(self.device)
        output, loss = self.model.loss(data, target)

        # Store extra batch information for calculation of the statistics
        batch_info['data'] = data
        batch_info['target'] = target
        batch_info['output'] = output
        batch_info['loss'] = loss

        return loss