How to make tracer account for arbitrary restrictions such as making a particular dim divisible by 4?

[saravanabalagi]

In my model, I check for shape during the forward function with something like:

if x.shape[1] % 4 != 0:
    raise ValueError(
        "Cannot use sin/cos positional encoding with "
        "odd dimension (got dim={:d})".format(x.shape[1])
    )

It is ok to prune as long as x.shape[1] is divisible by 4, but since this information is not available to the tracer, it prunes to match the specified pruning ratio. How can I make sure this restriction is accounted for during pruning?

[saravanabalagi]

Related to #38. Although it is entirely different with regards to its objective, the solution can in part solve this one as well. But the param round_to will round it for all layers, if I am not wrong. Is there another/better way to solve this problem?