From 6a8972ef9aad05ddf92564eca4e554ba9db8e1d9 Mon Sep 17 00:00:00 2001
From: Ethan Harris <ethanwharris@gmail.com>
Date: Wed, 1 Dec 2021 21:32:14 +0000
Subject: [PATCH] Fix

---
 flash/image/segmentation/data.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flash/image/segmentation/data.py b/flash/image/segmentation/data.py
index a747bfbe1b..6d62878cb5 100644
--- 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