Define input size of fully connected layer

[khalilDF0]

Hello All, I'm having issues in defining the input size for fully connected layer after a choice of two convolution layers like the code bellow shows:
The convlution block:

self.pool2 =LayerChoice([
nn.AvgPool2d(kernel_size=2, stride=2),
nn.AvgPool2d(kernel_size=4, stride=2)
], label='pool2')

the layer block
feature = nni.choice('feature', [64, 128, 256])
self.fc1 = MutableLinear(50176, feature)

so the issue is defining the inputs of the mutable linear, since the inputs changes after every trail based on the search space.