If the function accepts arguments in any precision, and doesn’t even care if they match, you probably don’t need to do anything.
If you’re wrapping the op with a torch.autograd.Function:
If the kernel needs matching input dtypes, but those dtypes can be anything as long as they match, the approach you describe is probably the best.
If the cuda kernel needs a particular dtype, you can do it a bit more conveniently 9.
If you’re registering the op with torch on the C++ side, see the autocast section of the dispatch registration tutorial 6.
autocast(enabled=False) does not affect apex amp. torch.cuda.amp is meant as the permanent and complete replacement, so you should prefer torch.cuda.amp.