with tf.variable_scope("conv_f"):
f = conv(x, 1, ch // 8, 1) # [bs, h, w, c']
with tf.variable_scope("conv_g"):
g = conv(x, 1, ch // 8, 1) # [bs, h, w, c']
with tf.variable_scope("conv_h"):
h = conv(x, 1, ch, 1) # [bs, h, w, c]
# N = h * w
s = tf.matmul(hw_flatten(g), hw_flatten(f), transpose_b=True) # # [bs, N, N]
beta = tf.nn.softmax(s) # attention map
o = tf.matmul(beta, hw_flatten(h)) # [bs, N, C]
gamma = tf.get_variable("gamma", [1], initializer=tf.constant_initializer(0.0))
o = tf.reshape(o, shape=[batch_size, height, width, num_channels // 2]) # [bs, h, w, C]
o = conv(o, 1, ch, 1)
x = gamma * o + x
return x
i got this error " ValueError: Cannot reshape a tensor with 16777216 elements to shape [32,64,64,64] (8388608 elements) for 'generator/encoder_5/Reshape_3' (op: 'Reshape') with input shapes: [32,4096,128], [4] and with input tensors computed as partial shapes: input[1] = [32,64,64,64].