Merge pull request MPolaris#9 from MPolaris/dev

Expand update
Alter-Y · Sep 23, 2022 · 22e2a9f · 22e2a9f
2 parents dee84a0 + 981c07a
commit 22e2a9f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/layers/deformation_layers.py b/layers/deformation_layers.py
@@ -144,6 +144,8 @@ def __call__(self, inputs):
         for i in range(len(self.shape)):
             if int(self.shape[i]//inputs.shape[i]) > 1:
                 inputs = tf.repeat(inputs, repeats=int(self.shape[i]//inputs.shape[i]), axis=i)
+            elif self.shape[i] < inputs.shape[i] and self.shape[i] != 1:
+                inputs = tf.repeat(inputs, repeats=int(self.shape[i]), axis=i)
         return inputs
 
 @OPERATOR.register_operator("Unsqueeze")