Update following comments

Lightning-Universe · Nov 4, 2021 · 3a50eda · 3a50eda
1 parent a657b20
commit 3a50eda
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/flash/core/model.py b/flash/core/model.py
@@ -624,10 +624,10 @@ def output(self) -> Optional[Output]:
 
     @torch.jit.unused
     @output.setter
-    def output(self, serializer: Union[Output, Mapping[str, Output]]):
-        if isinstance(serializer, Mapping):
-            serializer = OutputMapping(serializer)
-        self._output = serializer
+    def output(self, output: Union[Output, Mapping[str, Output]]):
+        if isinstance(output, Mapping):
+            output = OutputMapping(output)
+        self._output = output
 
     def build_data_pipeline(
         self,