Update capsule_test.py

earlzz · Feb 6, 2018 · 8c14be4 · 8c14be4
1 parent 82314b8
commit 8c14be4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/capsule_test.py b/capsule_test.py
@@ -81,7 +81,7 @@
 cnn = Conv2D(128, (3, 3), activation='relu')(cnn)
 cnn = Reshape((-1, 128))(cnn)
 capsule = Capsule(10, 16, 3, True)(cnn)
-output = Lambda(lambda x: K.sqrt(K.sum(K.square(x), 2)))(capsule)
+output = Lambda(lambda x: K.sqrt(K.sum(K.square(x), 2)), output_shape=(10,))(capsule)
 
 model = Model(inputs=input_image, outputs=output)
 model.compile(loss=lambda y_true,y_pred: y_true*K.relu(0.9-y_pred)**2 + 0.25*(1-y_true)*K.relu(y_pred-0.1)**2,