_phase_shift(I, r) error while using batch size 1

[Patryk Chrabaszcz]

I think there might be a problem with function _phase_shift(I,r). My code was crashing when image was resized from 1x1 to 2x2 . Probably it is because of some kind of dimension mismatch.
I changed it to :

def _phase_shift(I, r):
# Helper function with main phase shift operation
bsize, a, b, c = I.get_shape().as_list()
X = tf.reshape(I, (bsize, a, b, r, r))
X = tf.transpose(X, (0, 1, 2, 4, 3)) # bsize, a, b, 1, 1
X = tf.split(1, a, X) # a, [bsize, b, r, r]
X = tf.concat(2, [tf.squeeze(x, squeeze_dims=1) for x in X]) # bsize, b, ar, r
X = tf.split(1, b, X) # b, [bsize, ar, r]
X = tf.concat(2, [tf.squeeze(x, squeeze_dims=1) for x in X]) #
bsize, ar, br
return tf.reshape(X, (bsize, ar, br, 1))

Its just adding squeeze_dims=1 . Now it seems to work with batch size = 1

[Patryk Chrabaszcz]

Oh it was resolved.