Fixed a remnant hardcoding from 2D init

NifTK · Oct 24, 2019 · abd8e70 · abd8e70
1 parent 72437a2
commit abd8e70
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/niftynet/layer/subpixel.py b/niftynet/layer/subpixel.py
@@ -182,13 +182,13 @@ def __call__(self, shape, dtype, partition_info=None):
  return self.initializer(shape)
 
  # Initializing W0 (enough kernels for one output channel)
- new_shape = shape[:-1] + [shape[-1] // (self.upsample_factor ** 2)]
+ new_shape = shape[:-1] + [shape[-1] // (self.upsample_factor ** (len(shape)-2))]
  x = self.initializer(new_shape, dtype, partition_info)
 
  # Repeat the elements along the output dimension
  x = tf.keras.backend.repeat_elements(
  x=x,
- rep=self.upsample_factor ** 2,
+ rep=self.upsample_factor ** len(shape)-2,
  axis=-1
  )