ENSEEIHT/projet-long
1
0
Fork 0
Code Issues Pull requests Packages Projects Releases Wiki Activity

Merge pull request #188 from laclouis5/batchSizeGreaterThanOneFix

Browse source 
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
parent aa9defe6a5 416e076dbc
commit fb5f60fb31
1 changed files with 4 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
utils/dataset.py
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)
}