'Helper function for fit which performs training on a dataset for one epoch.
Arguments:
dataset (NervanaDataIterator): Dataset iterator to perform fit on'

def _epoch_fit(self, dataset, callbacks):

    epoch = self.epoch_index
    self.total_cost[:] = 0
    last_gen_iter = self.gen_iter
    (z, y_temp) = (self.zbuf, self.ybuf)
    for (mb_idx, (x, _)) in enumerate(dataset):
        callbacks.on_minibatch_begin(epoch, mb_idx)
        self.be.begin(Block.minibatch, mb_idx)
        if self.wgan_param_clamp:
            self.clip_param_in_layers(self.layers.discriminator.layers_to_optimize, self.wgan_param_clamp)
        self.fill_noise(z, normal=(self.noise_type == 'normal'))
        Gz = self.fprop_gen(z)
        y_noise = self.fprop_dis(Gz)
        y_temp[:] = y_noise
        delta_noise = self.cost.costfunc.bprop_noise(y_noise)
        self.bprop_dis(delta_noise)
        self.layers.discriminator.set_acc_on(True)
        y_data = self.fprop_dis(x)
        delta_data = self.cost.costfunc.bprop_data(y_data)
        self.bprop_dis(delta_data)
        self.optimizer.optimize(self.layers.discriminator.layers_to_optimize, epoch=epoch)
        self.layers.discriminator.set_acc_on(False)
        self.cost_dis[:] = self.cost.get_cost(y_data, y_temp, cost_type='dis')
        if (self.current_batch == (self.last_gen_batch + self.get_k(self.gen_iter))):
            self.fill_noise(z, normal=(self.noise_type == 'normal'))
            Gz = self.fprop_gen(z)
            y_temp[:] = y_data
            y_noise = self.fprop_dis(Gz)
            delta_noise = self.cost.costfunc.bprop_generator(y_noise)
            delta_dis = self.bprop_dis(delta_noise)
            self.bprop_gen(delta_dis)
            self.optimizer.optimize(self.layers.generator.layers_to_optimize, epoch=epoch)
            self.cost_dis[:] = self.cost.get_cost(y_temp, y_noise, cost_type='dis')
            self.total_cost[:] = (self.total_cost + self.cost_dis)
            self.last_gen_batch = self.current_batch
            self.gen_iter += 1
        self.be.end(Block.minibatch, mb_idx)
        callbacks.on_minibatch_end(epoch, mb_idx)
        self.current_batch += 1
    assert (self.gen_iter > last_gen_iter), 'at    least    one    generator    iteration    is    required    for    total    cost    estimation    in    this    epoch'
    self.total_cost[:] = (self.total_cost / (self.gen_iter - last_gen_iter))
    (self.data_batch, self.noise_batch) = (x, self.fprop_gen(self.z0))
