Fix keras shape inference for lambda outputs (dmlc#84)

joeddav · Sep 16, 2017 · a21823b · a21823b
1 parent 928d1a5
commit a21823b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/keras/layers/core.py b/keras/layers/core.py
@@ -599,8 +599,8 @@ def __init__(self, function, output_shape=None, arguments=None, **kwargs):
 
     def get_output_shape_for(self, input_shape):
         if self._output_shape is None:
-            # With TensorFlow, we can infer the output shape directly:
-            if K.backend() == 'tensorflow':
+            # With TensorFlow and MXNet, we can infer the output shape directly:
+            if K.backend() == 'tensorflow' or 'mxnet':
                 if isinstance(input_shape, list):
                     xs = [K.placeholder(shape=shape) for shape in input_shape]
                     x = self.call(xs)