tests

Lightning-Universe · Aug 17, 2021 · 4931cd7 · 4931cd7
1 parent 3c7e2db
commit 4931cd7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/flash/image/segmentation/data.py b/flash/image/segmentation/data.py
@@ -62,6 +62,7 @@
     plt = None
 
 if _TORCHVISION_AVAILABLE:
+    import torchvision
     import torchvision.transforms.functional as FT
     from torchvision.datasets.folder import default_loader, has_file_allowed_extension, IMG_EXTENSIONS
 else:
@@ -149,7 +150,7 @@ def load_sample(self, sample: Mapping[str, Any]) -> Mapping[str, Union[torch.Ten
 
         # load images directly to torch tensors
         img: torch.Tensor = FT.to_tensor(default_loader(img_path))  # CxHxW
-        img_labels: torch.Tensor = FT.to_tensor(default_loader(img_labels_path))  # CxHxW
+        img_labels: torch.Tensor = torchvision.io.read_image(img_labels_path)  # CxHxW
         img_labels = img_labels[0]  # HxW
 
         sample[DefaultDataKeys.INPUT] = img.float()