From 2ed9194d50159745e4ad58a1a6b5c4223f2f3bbf Mon Sep 17 00:00:00 2001 From: Ethan Harris Date: Mon, 13 Dec 2021 14:14:11 +0000 Subject: [PATCH] Fixes --- flash/audio/speech_recognition/output_transform.py | 2 +- flash/image/segmentation/model.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/flash/audio/speech_recognition/output_transform.py b/flash/audio/speech_recognition/output_transform.py index d5f65a2df0..56f282bed5 100644 --- a/flash/audio/speech_recognition/output_transform.py +++ b/flash/audio/speech_recognition/output_transform.py @@ -62,7 +62,7 @@ def per_batch_transform(self, batch: Any) -> Any: def __getstate__(self): # TODO: Find out why this is being pickled state = self.__dict__.copy() - state.pop("_tokenizer") + state.pop("_tokenizer", None) return state def __setstate__(self, state): diff --git a/flash/image/segmentation/model.py b/flash/image/segmentation/model.py index 84d3ed566f..c27b589958 100644 --- a/flash/image/segmentation/model.py +++ b/flash/image/segmentation/model.py @@ -34,9 +34,9 @@ OUTPUT_TRANSFORM_TYPE, OUTPUT_TYPE, ) -from flash.image.data import ImageDeserializer from flash.image.segmentation.backbones import SEMANTIC_SEGMENTATION_BACKBONES from flash.image.segmentation.heads import SEMANTIC_SEGMENTATION_HEADS +from flash.image.segmentation.input import SemanticSegmentationDeserializer from flash.image.segmentation.output import SegmentationLabelsOutput from flash.image.segmentation.transforms import SemanticSegmentationInputTransform @@ -184,7 +184,7 @@ def serve( host: str = "127.0.0.1", port: int = 8000, sanity_check: bool = True, - input_cls: Optional[Type[ServeInput]] = ImageDeserializer, + input_cls: Optional[Type[ServeInput]] = SemanticSegmentationDeserializer, transform: INPUT_TRANSFORM_TYPE = SemanticSegmentationInputTransform, transform_kwargs: Optional[Dict] = None, ) -> Composition: