def train(epoch):
    train_sampler.set_epoch(epoch)
    model.train()
    losses = AverageMeter()
    top1 = AverageMeter()
    global best_pred, acclist_train
    tic = time.time()
    for (batch_idx, (data, target)) in enumerate(train_loader):
        scheduler(optimizer, batch_idx, epoch, best_pred)
        if (not args.mixup):
            (data, target) = (data.cuda(args.gpu), target.cuda(args.gpu))
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        if args.amp:
            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            loss.backward()
        optimizer.step()
        if (not args.mixup):
            acc1 = accuracy(output, target, topk=(1,))
            top1.update(acc1[0], data.size(0))
        losses.update(loss.item(), data.size(0))
        if (((batch_idx % 100) == 0) and (args.gpu == 0)):
            iter_per_sec = ((100.0 / (time.time() - tic)) if (batch_idx != 0) else (1.0 / (time.time() - tic)))
            tic = time.time()
            if args.mixup:
                print('Epoch: {}, Iter: {}, Speed: {:.3f} iter/sec, Train loss: {:.3f}'.format(epoch, batch_idx, iter_per_sec, losses.avg.item()))
            else:
                print('Epoch: {}, Iter: {}, Speed: {:.3f} iter/sec, Top1: {:.3f}'.format(epoch, batch_idx, iter_per_sec, top1.avg.item()))
    acclist_train += [top1.avg]
def train(epoch):
    train_sampler.set_epoch(epoch)
    model.train()
    losses = AverageMeter()
    top1 = AverageMeter()
    global best_pred, acclist_train
    tic = time.time()
    for (batch_idx, (data, target)) in enumerate(train_loader):
        scheduler(optimizer, batch_idx, epoch, best_pred)
        if (not args.mixup):
            (data, target) = (data.cuda(args.gpu), target.cuda(args.gpu))
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        if args.amp:
            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            loss.backward()
        optimizer.step()
        if (not args.mixup):
            acc1 = accuracy(output, target, topk=(1,))
            top1.update(acc1[0], data.size(0))
        losses.update(loss.item(), data.size(0))
        if (((batch_idx % 100) == 0) and (args.gpu == 0)):
            iter_per_sec = ((100.0 / (time.time() - tic)) if (batch_idx != 0) else (1.0 / (time.time() - tic)))
            tic = time.time()
            if args.mixup:
                print('Epoch: {}, Iter: {}, Speed: {:.3f} iter/sec, Train loss: {:.3f}'.format(epoch, batch_idx, iter_per_sec, losses.avg.item()))
            else:
                print('Epoch: {}, Iter: {}, Speed: {:.3f} iter/sec, Top1: {:.3f}'.format(epoch, batch_idx, iter_per_sec, top1.avg.item()))
    acclist_train += [top1.avg]