im_data = torch.FloatTensor(1).cuda()
....
for step in range(...):
batch_data = next(train_iter):
im_data.data.resize_(batch_data[0].shape).copy_(batch_data[0])
scores = net(im_data)
....
1 Like