The the Autocast and Custom Autograd Functions section use cases for custom methods are described. In your use case, since you don’t want to allow the inputs to be transformed to float16, you could disable autocast via a nested context manager:

In all cases, if you’re importing the function and can’t alter its definition, a safe fallback is to disable autocast and force execution in float32 ( or dtype) at any points of use where errors occur:

with autocast():
    ...
    with autocast(enabled=False):
        output = imported_function(input1.float(), input2.float())