Загрузчик данных Pytorch перемешивает пример внутри пакета

Question

У меня есть следующий Dataloader, который я использую для автоматического кодера в Pytorch

class DataLoader(data.DataLoader):
    def __init__(self,  *args, **kwargs):
        super(DataLoader, self).__init__(*args, **kwargs)
        self.iterator = self.__iter__()

    def __next__(self):
        try:
            return next(self.iterator)
        except:
            self.iterator = self.__iter__()
            return next(self.iterator)

При выполнении вычисления функции потерь я хочу отображать прогнозы в виде изображений для первого файла в моем наборе данных, но даже если у меня есть только одна партия из 2 изображений для обучения, эти изображения обмениваются, поэтому мой код:

class CrossEntropyLoss2d(nn.Module):
    def __init__(self, weight=None):
        super(CrossEntropyLoss2d, self).__init__()

    def forward(self, pred, target, weight=None):
        pred = pred.clamp(min=1e-16)

        plt.clf()
        predpic = (np.uint8(np.argmax(pred.data.cpu().numpy(), 1)))
        plt.imshow(predpic[0]) ### here I want only the first picture! 
        plt.title('Example prediction')
        plt.savefig(os.path.join(results_path, "exampleprediction.png"))
        plt.show()

        ##### COMPUTE LOSS
        loss = -(pred.log() * target)
        loss = loss.sum(1)
        if weight is not None:
            loss = loss * weight
        loss = loss.mean()
        return loss 

не работает.

Как я могу это исправить? Спасибо!