Fix implicit cast in custom_function (#46445)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/46445
Fix an instance in which a truncated integer prevents downstream type safety checks.
Test Plan: I'm not sure what's appropriate here.
Reviewed By: albanD
Differential Revision: D24339292
fbshipit-source-id: 15748ec64446344ff1a8344005385906d3484d7c
