bug-fix cifar10_cnn_capsule.py missing K.sum() (keras-team#9520)

Just add K.sum() around the margin_loss
ferhanture · Mar 1, 2018 · 500401b · 500401b
1 parent 98ffe08
commit 500401b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/cifar10_cnn_capsule.py b/examples/cifar10_cnn_capsule.py
@@ -44,8 +44,8 @@ def softmax(x, axis=-1):
 # define the margin loss like hinge loss
 def margin_loss(y_true, y_pred):
     lamb, margin = 0.5, 0.1
-    return y_true * K.square(K.relu(1 - margin - y_pred)) + lamb * (
-        1 - y_true) * K.square(K.relu(y_pred - margin))
+    return K.sum(y_true * K.square(K.relu(1 - margin - y_pred)) + lamb * (
+        1 - y_true) * K.square(K.relu(y_pred - margin)), axis=-1)
 
 
 class Capsule(Layer):