Moodhi
Moodhi

Reputation: 45

How can I send a Data Loader to the GPU in Google Colab?

I have two data loaders and I am trying to send them to the GPU using .to(device) but this is not working. This is the code I am using:

# to create a batch iterator

class MyData(Dataset):
    def __init__(self, X, y):
        self.data = X
        self.target = y
        # TODO: convert this into torch code is possible
        self.length = [ np.sum(1 - np.equal(x, 0)) for x in X]
        
    def __getitem__(self, index):
        x = self.data[index]
        y = self.target[index]
        x_len = self.length[index]
        xx_len = torch.tensor(x_len)
        return {"src": x, "trg": y, "x_len": xx_len}
    
    def __len__(self):
        return len(self.data)


dataset = DataLoader(train_dataset, batch_size = BATCH_SIZE, 
                     drop_last=True,
                     shuffle=True)
test_Dataset= DataLoader(val_dataset, batch_size = BATCH_SIZE, 
                     drop_last=True,
                     shuffle=True)

I also tried to use pin_memory = True but this is not working too.

Upvotes: 1

Views: 1162

Answers (1)

user14518353
user14518353

Reputation:

You do not move the dataloader to the GPU. Instead, create batch tensors that store the data and then move these tensors to the GPU.

train_dataloader = DataLoader(MyData, batch_size=BS)
...
def train(nn, optim, train_dataloader, etc...): 
    for batch in train_dataloader:
        batch = batch.to('cuda')
    ...

Upvotes: 2

Related Questions