can you describe the whole flow? Are you trying to impose the limit on the kernel level or when people train the model?
If you need a kernel that imposes some limits on the int32 value, then I think the best thing to do is to reimplement the kernel (possibly by calling fbgemm implementations if you need high performance: GitHub - pytorch/FBGEMM: FB (Facebook) + GEMM (General Matrix-Matrix Multiplication) - https://code.fb.com/ml-applications/fbgemm/)