Could not get batch from DataLoader

Viewed 29

I'm trying to implement a model using lstm to generate sonnects in pytorch. When I tested the dataLoader, It took a lot of time and could not return data as expected. I review it and don't know where error locates. Please help me, here is the code(more detail in(https://paste.ubuntu.com/p/vspS3msNVW/)) (10th cell)

class SonnetDataset(Dataset):
    def __init__(self, sonnet_in_ids: list, vocab: list, max_seq_length: int):
        super().__init__()
        self.data = sonnet_in_ids
        self.vocab = vocab
        self.vocab_size = len(vocab)
        self.pad_id = self.vocab.index('<PAD>')
        self.start_id = self.vocab.index('<START>')
        self.end_id = self.vocab.index('<END>')
        self.max_seq_length = max_seq_length + 2
        print('init successfully')

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index) -> torch.LongTensor:

        print('get item')
        x = self.data[index]
        x = [self.start_id] + x + [self.end_id]
        # padding 
        x += [self.pad_id] * (self.max_seq_length - len(x))
        x = torch.LongTensor(x)
        print(x)
        return x
batch_size = 4
train_set = SonnetDataset(sonnet_in_ids=sonnets_in_ids, vocab=vocab, max_seq_length=max_length)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=4, drop_last=True)


next(iter(train_loader))

Could not get data

0 Answers
Related