Allow None in Keras loss multitask loss (#179)

When training multi output models Keras allows to use `loss=None` for some branches. Adapt typings to support this.
larq · Aug 12, 2020 · 230b455 · 230b455
1 parent 61b1162
commit 230b455
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/zookeeper/tf/experiment.py b/zookeeper/tf/experiment.py
@@ -22,6 +22,6 @@ class Experiment:
     epochs: int = Field()
     batch_size: int = Field()
     loss: Optional[
-        Union[Sequence[Union[keras.losses.Loss, str]], keras.losses.Loss, str]
+        Union[Sequence[Union[keras.losses.Loss, str, None]], keras.losses.Loss, str]
     ] = Field()
     optimizer: Union[keras.optimizers.Optimizer, str] = Field()