@ptrblck
actually coming back to this, because explicit autocasting does not sit well while converting the model to torchscript and complains
torch.jit.frontend.UnsupportedNodeError: function definitions aren't supported:

def foo(
    x: torch.Tensor, y: torch.Tensor
) -> torch.Tensor:
    with autocast(enabled=False):
       ..........
    return z

how can we make autocast context manager and torchscript conversion co-exist ?