From 416e076dbcb312051bf87b46cdda9e803244163f Mon Sep 17 00:00:00 2001
From: Louis Lac <lac.louis5@gmail.com>
Date: Fri, 5 Jun 2020 19:00:04 +0200
Subject: [PATCH] Conversion to FloatTensor

Former-commit-id: cc44eebf19f5d98fbd1ca800608ae0b38d998b54
---
 utils/dataset.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/utils/dataset.py b/utils/dataset.py
index bfba1a6..3afeca8 100644
--- 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)
+        }