Hi, my training codes are as follows
model.train()
for batch_idx, (input, target) in enumerate(loader):
    target = target.cuda(non_blocking=True)
    optimizer.zero_grad()
    width_mult_list = []
    # first do max_width and max_resolution
    max_width = FLAGS.width_mult_range[1]
    model.apply(lambda m: setattr(m, 'width_mult', max_width))
    max_output = model(input)
    loss = torch.mean(criterion(max_output, target))
    loss.backward()
    max_output = max_output.detach()
    # do other widths and resolution
    min_width = FLAGS.width_mult_range[0]
    width_mult_list = [min_width]
    for i in range(2):
        sampled_width = random.uniform(FLAGS.width_mult_range[0], FLAGS.width_mult_range[1])
        width_mult_list.append(sampled_width)
    for width_mult in sorted(width_mult_list, reverse=True):
        model.apply(
            lambda m: setattr(m, 'width_mult', width_mult))
        resolution = FLAGS.resolution_list[random.randint(0,3)]
        subinput = F.interpolate(input, (resolution, resolution), mode='bilinear', align_corners=True)
        output = model(subinput)
        loss = torch.nn.KLDivLoss(reduction='batchmean')(F.log_softmax(output, dim=1), F.softmax(max_output, dim=1))
        loss.backward()
    optimizer.step()