Solved by Crazyai in post #2 
              
              
                I am a freshman to FX. Maybe fx.wrap can help you. 
@fx.wrap
def get_size(x, y):
    return torch.Size([x, y])

class Example (nn.Module):
    def __init__(self) -> None:
        super(Example, self).__init__()

    def forward(self, input):

        batch_size = input.size(0)
        ch_size = inpu…