@hm2092 thank you for sharing your error - I think I am experiencing something very similar where the cuDNN message only happens for batch size > 1. Could you explain more about how you debugged the location of the “out of range” error? My DataLoader + custom collate_fn is fairly straightforward and there are no issues preparing the batches
On a separate note, why did PyTorch only implement a Dataset class for COCO, but not a DataLoader and/or collate function? It seems this is standard enough that plenty of other people are using it…
from spellchecker import SpellChecker
import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.datasets import CocoCaptions

class CocoCollate:
    """Custom collate function for COCO captions dataset

    :param dict stoi: Mapping defining integer values for each word in the vocabulary
    """

    def __init__(self, stoi, padding_idx=-1):
        self.stoi = stoi
        self.unk = len(self.stoi)
        self.padding_idx = padding_idx

    def __call__(self, batch):
        """Randomly sample caption and numericalize using mapping

        :param tuple(torch.tensor, [[str]]) batch: 3-channel image data and nested
            lists of potential captions (5 per sample)
        :return tuple(torch.tensor): Batched images and numericalized captions
        """
        images, captions = zip(*batch)
        spell = SpellChecker(distance=1)
        selected = [[word if word in self.stoi else spell.correction(word)
                     for word in re.findall(r"[<>\w']+|[.,!?;]", pick)]
                    for pick in [f'<s> {c[random.randrange(5)].lower()} </s>' for c in captions]]
        numericalized = [torch.tensor([self.stoi.get(word, self.unk) for word in s]) for s in selected]
        packed = nn.utils.rnn.pad_sequence(numericalized, batch_first=True, padding_value=self.padding_idx)
        return torch.stack(images), packed

transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
trainset = datasets.CocoCaptions('/path/to/images', '/path/to/captions', transform)
vocabulary = [ ... ]
stoi = {word: idx for idx, word in enumerate(vocabulary)}
loader = DataLoader(trainset, batch_size=2, collate_fn=CocoCollate(stoi))