How to use inbuilt `int` in torch.fx

I’m trying to symbolically trace these lines:

B, T, _, C = x.shape
grid_size = int(math.sqrt(T))

The problem is that int can’t handle a Proxy as an input. Is there any way I can customize the tracer to handle this particular use case?