dataset = ListDataset(train_path, augment=True, multiscale=opt.multiscale_training)
dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=opt.batch_size,
    shuffle=True,
    num_workers=opt.n_cpu,
    pin_memory=True,
    collate_fn=dataset.collate_fn
)

optimizer = torch.optim.SGD(model.parameters(), lr=opt.lr, momentum=0.9)

        imgs = imgs.to(device)
        targets = targets.to(device)

        loss, outputs = model(imgs, targets)

        optimizer.zero_grad()
        loss.backward()

        updateBN(sr_flag, model.module_list, opt.s, prune_idx)
        #BNOptimizer.updateBN(sr_flag, model.module_list, opt.s, prune_idx)

        optimizer.step()