Dataloader adds additional unwanted dimension between __getitem__ and enumerate(dataloader)

Hello everybody,
I have a dataloader, with a __getitem__ function, that reads as follows:

    def __getitem__(self, index):
        
        # pre-processing annotation
        with torch.no_grad():

            annotation = torch.zeros((self.image_size[0], self.image_size[1]))
            points = np.loadtxt(os.path.join(self.annotation_root, self.annotations[index]), delimiter=",")
            image = Image.open(os.path.join(self.image_root, self.images[index]))
            img_size = image.size

            # converting relative coordinates to absolute ones

            if len(points.shape) < 2: 
                points = np.expand_dims(points, axis = 0)
            points = points[:,0:2]
            points[:,0] = np.rint(points[:,0]*img_size[0])
            points[:,1] = np.rint(points[:,1]*img_size[1])


            annotation[points[:,0], points[:,1]] = 1.0
            count = points.shape[0]

            # annotation is padded by receptive field size by default
            annotation = annotation.view((1,1,*annotation.size()))
            sum_annotation = F.conv2d(input = annotation,
                                      weight = self.weights,
                                      stride = self.stride,
                                      padding = self.patch_size)

        print("sum_annotation.size(): {}".format(sum_annotation.size()))

        return self.make_tensor(image), sum_annotation, count

nothing special about it, just some pre-processing for my annotations. Here, the print returns sum_annotation.size(): torch.Size([1, 1, 921, 1041]). However, doing

    for epoch in range(args.epochs):
        loss = 0
        for idx, elt in enumerate(train_dataloader):
            input, target, _ = elt
            print("target.size(): {}".format(target.size()))

gets me a strange print output: target.size(): torch.Size([1, 1, 1, 921, 1041]). Somewhere, an unwanted dimension gets added. I could do something ugly using tensor.view to fix this, but I would like to know if anybody has ever encountered this issue before.
Any help would be much appreciated :slight_smile:

The DataLoader will add the batch dimension to the sample in dim0. Based on your description it seems that you are using batch_size=1 which is why the additional dimension with a size of 1 is added.
I don’t know what the other dimension represent, but if you are manually already trying to add a batch dimension, just remove it and let the DataLoader handle it.

of course, that’s obvious now ! thanks, I had a brainfart … I needed the batch dimension for a pre-processing step involving F.conv2d, but I forgot to remove it before returning this tensor.