Fix

Lightning-Universe · Dec 1, 2021 · 6a8972e · 6a8972e
1 parent 4958df0
commit 6a8972e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flash/image/segmentation/data.py b/flash/image/segmentation/data.py
@@ -194,7 +194,7 @@ class SemanticSegmentationDeserializer(ImageDeserializer):
     def deserialize(self, data: str) -> Dict[str, Any]:
         result = super().deserialize(data)
         result[DataKeys.INPUT] = FT.to_tensor(result[DataKeys.INPUT])
-        result[DataKeys.METADATA] = {"size": result[DataKeys.INPUT].shape}
+        result[DataKeys.METADATA] = {"size": result[DataKeys.INPUT].shape[-2:]}
         return result