如何将数据加载器发送到 Google Colab 中的 GPU?
pytorch 491
原文标题 :How can I send a Data Loader to the GPU in Google Colab?
我有两个数据加载器,我正在尝试使用.to(device)
将它们发送到 GPU,但这不起作用。这是我正在使用的代码:
# 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)
我也尝试使用pin_memory = True
但这也不起作用。