Merge pull request #188 from laclouis5/batchSizeGreaterThanOneFix

Fixed Dataloader for batch size greater than one

Former-commit-id: 5419352796a0b221c2d8f9c47f2f7c1db2459c54
This commit is contained in:
milesial 2020-06-05 16:35:12 -07:00 committed by GitHub
commit fb5f60fb31

View file

@ -59,4 +59,7 @@ class BasicDataset(Dataset):
img = self.preprocess(img, self.scale) img = self.preprocess(img, self.scale)
mask = self.preprocess(mask, self.scale) mask = self.preprocess(mask, self.scale)
return {'image': torch.from_numpy(img), 'mask': torch.from_numpy(mask)} return {
'image': torch.from_numpy(img).type(torch.FloatTensor),
'mask': torch.from_numpy(mask).type(torch.FloatTensor)
}