def run_train_epoch(model, optimizer, log_prior, dataloader, epoch, asr_decoder, trans_model, silence_ids, args):
    batch_time = utils.AverageMeter('Time', ':6.3f')
    losses = utils.AverageMeter('Loss', ':.4e')
    grad_norm = utils.AverageMeter('grad_norm', ':.4e')
    progress = utils.ProgressMeter(len(dataloader), batch_time, losses, grad_norm, prefix='Epoch: [{}]'.format(epoch))
    ce_criterion = nn.CrossEntropyLoss(ignore_index=(- 100), reduction='sum')
    if (args.criterion == 'mmi'):
        se_criterion = ops.MMIFunction.apply
    else:
        se_criterion = ops.sMBRFunction.apply
    end = time.time()
    for (i, batch) in enumerate(dataloader, 0):
        feat = batch['x']
        label = batch['y']
        num_frs = batch['num_frs']
        utt_ids = batch['utt_ids']
        aux = batch['aux']
        x = feat.to(th.float32)
        y = label.long()
        x = x.cuda()
        y = y.cuda()
        prediction = model(x)
        ce_loss = ce_criterion(prediction.view((- 1), prediction.shape[2]), y.view((- 1)))
        se_loss = 0.0
        for j in range(len(num_frs)):
            log_like_j = prediction[j, :, :]
            log_like_j = log_like_j[:num_frs[j], :]
            log_like_j = (log_like_j - log_prior)
            trans_id = th.from_numpy(aux[j][0][0].astype(int)).tolist()
            if (args.criterion == 'mmi'):
                se_loss += se_criterion(log_like_j, asr_decoder, trans_model, trans_id)
            else:
                se_loss += se_criterion(log_like_j, asr_decoder, trans_model, trans_id, args.criterion, silence_ids)
        loss = (se_loss.cuda() + (args.ce_ratio * ce_loss))
        optimizer.zero_grad()
        loss.backward()
        norm = nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
        optimizer.step()
        grad_norm.update(norm)
        tot_frs = np.array(num_frs).sum()
        losses.update((loss.item() / tot_frs))
        batch_time.update((time.time() - end))
        if ((hvd.rank() == 0) and ((i % args.save_freq) == 0)):
            checkpoint = {}
            checkpoint['model'] = model.state_dict()
            checkpoint['optimizer'] = optimizer.state_dict()
            output_file = (((args.exp_dir + '/model.se.') + str(i)) + '.tar')
            th.save(checkpoint, output_file)
        if ((hvd.rank() == 0) and ((i % args.print_freq) == 0)):
            progress.print(i)
def run_train_epoch(model, optimizer, log_prior, dataloader, epoch, asr_decoder, trans_model, silence_ids, args):
    batch_time = utils.AverageMeter('Time', ':6.3f')
    losses = utils.AverageMeter('Loss', ':.4e')
    grad_norm = utils.AverageMeter('grad_norm', ':.4e')
    progress = utils.ProgressMeter(len(dataloader), batch_time, losses, grad_norm, prefix='Epoch: [{}]'.format(epoch))
    ce_criterion = nn.CrossEntropyLoss(ignore_index=(- 100), reduction='sum')
    if (args.criterion == 'mmi'):
        se_criterion = ops.MMIFunction.apply
    else:
        se_criterion = ops.sMBRFunction.apply
    end = time.time()
    for (i, batch) in enumerate(dataloader, 0):
        feat = batch['x']
        label = batch['y']
        num_frs = batch['num_frs']
        utt_ids = batch['utt_ids']
        aux = batch['aux']
        x = feat.to(th.float32)
        y = label.long()
        x = x.cuda()
        y = y.cuda()
        prediction = model(x)
        ce_loss = ce_criterion(prediction.view((- 1), prediction.shape[2]), y.view((- 1)))
        se_loss = 0.0
        for j in range(len(num_frs)):
            log_like_j = prediction[j, :, :]
            log_like_j = log_like_j[:num_frs[j], :]
            log_like_j = (log_like_j - log_prior)
            trans_id = th.from_numpy(aux[j][0][0].astype(int)).tolist()
            if (args.criterion == 'mmi'):
                se_loss += se_criterion(log_like_j, asr_decoder, trans_model, trans_id)
            else:
                se_loss += se_criterion(log_like_j, asr_decoder, trans_model, trans_id, args.criterion, silence_ids)
        loss = (se_loss.cuda() + (args.ce_ratio * ce_loss))
        optimizer.zero_grad()
        loss.backward()
        norm = nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
        optimizer.step()
        grad_norm.update(norm)
        tot_frs = np.array(num_frs).sum()
        losses.update((loss.item() / tot_frs))
        batch_time.update((time.time() - end))
        if ((hvd.rank() == 0) and ((i % args.save_freq) == 0)):
            checkpoint = {}
            checkpoint['model'] = model.state_dict()
            checkpoint['optimizer'] = optimizer.state_dict()
            output_file = (((args.exp_dir + '/model.se.') + str(i)) + '.tar')
            th.save(checkpoint, output_file)
        if ((hvd.rank() == 0) and ((i % args.print_freq) == 0)):
            progress.print(i)