update

Lightning-Universe · May 12, 2021 · 8349d40 · 8349d40
1 parent 1da7d89
commit 8349d40
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/vision/segmentation/test_serialization.py b/tests/vision/segmentation/test_serialization.py
@@ -1,6 +1,7 @@
 import pytest
 import torch
 
+from flash.data.data_source import DefaultDataKeys
 from flash.vision.segmentation.serialization import SegmentationLabels
 
 
@@ -30,7 +31,7 @@ def test_serialize(self):
         sample[1, 1, 2] = 1  # add peak in class 2
         sample[3, 0, 1] = 1  # add peak in class 4
 
-        classes = serial.serialize(sample)
+        classes = serial.serialize({DefaultDataKeys.PREDS: sample})
         assert classes[1, 2] == 1
         assert classes[0, 1] == 3