Hi, thanks for your reply!
Here is the forward pass part in the script. I don’t know whether it is helpful.
def training(self, epoch):
    train_loss = 0.0
    self.model.train()
    tbar = tqdm(self.trainloader)
    for i, (image, target) in enumerate(tbar):
        image, target = image.to(self.device), target.to(self.device)
        self.scheduler(self.optimizer, i, epoch, self.best_pred)
        self.optimizer.zero_grad()
        outputs = self.model(image)

        # Multi-size training
        iter_loss = 0
        for logit in outputs:
            _, _, H, W = logit.shape
            labels_ = utils.resize_labels(target, size=(H, W))
            iter_loss += self.criterion(logit.cuda(), labels_.cuda())

        torch.cuda.empty_cache()
        iter_loss.backward()
        self.optimizer.step()
        train_loss += iter_loss.item()
        tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))

And I’ve not applied gather function (like torch.nn.parallel.scatter_gather.gather) explicitly.
thanks!