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 = input.size(1)
        # return torch.Size([batch_size, ch_size])
        return get_size(batch_size, ch_size)