Merge pull request #188 from laclouis5/batchSizeGreaterThanOneFix

Fixed Dataloader for batch size greater than one Former-commit-id: 5419352796a0b221c2d8f9c47f2f7c1db2459c54
2020-06-05 16:35:12 -07:00 · 2020-06-05 16:35:12 -07:00 · fb5f60fb31
parent aa9defe6a5 416e076dbc
commit fb5f60fb31
1 changed files with 4 additions and 1 deletions
--- a/utils/dataset.py
+++ b/utils/dataset.py
@ -59,4 +59,7 @@ class BasicDataset(Dataset):
        img = self.preprocess(img, 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)
+        }