Update Capsule_Keras.py

earlzz · Jan 24, 2018 · f9d8140 · f9d8140
1 parent 1e4bc16
commit f9d8140
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/Capsule_Keras.py b/Capsule_Keras.py
@@ -22,15 +22,15 @@ def build(self, input_shape):
  self.input_dim_capsule = input_shape[-1]
  if self.share_weights:
  self.W = self.add_weight(name='capsule_kernel',
- shape=(1, input_shape[-1],
+ shape=(1, self.input_dim_capsule,
  self.num_capsule * self.dim_capsule),
  initializer='glorot_uniform',
  trainable=True)
  else:
  self.input_num_capsule = input_shape[-2]
  self.W = self.add_weight(name='capsule_kernel',
  shape=(self.input_num_capsule,
- input_shape[-1],
+ self.input_dim_capsule,
  self.num_capsule * self.dim_capsule),
  initializer='glorot_uniform',
  trainable=True)