Hey @ptrblck, I’ve tried this for calculating the pixel-wise loss using the same SupConLoss as above.
def training(self, epoch):
        train_loss = 0.0
        self.model.train()
        tbar = tqdm(self.trainloader)
        for k, (image, target) in enumerate(tbar):
            self.scheduler(self.optimizer, k, epoch, self.best_pred)
            self.optimizer.zero_grad()
            image, target = image.cuda(), target.cuda()
            features = self.model(image)
            b,c,h,w = features.size()
            for i in range(h):
                for j in range(w):
                    f = features[:,:,i,j]          #[b,c]
                    f = f.unsqueeze(2)             #[b,c,1]
                    t = target[:,i,j]              #[b]
                    loss = self.criterion(f, t)
                    loss.sum().backward(retain_graph=True)
                    self.optimizer.step()
            return loss
            train_loss += loss.item()
            tbar.set_description('Train loss: %.3f' % (train_loss / (k + 1)))

My GPU utilization is at 95% but I’m not sure if the code is running or not. Can you double check it once. Thanks:)