Repository: pytorch
Issue: amp with `bf16`: backward happens in `f16` when using `@torch.cuda.amp.custom_bwd`
Labels: triaged,module: bfloat16,module: half,module: amp (automated mixed precision)

### 🐛 Describe the bug

I don't think it's expected, and it can be quite surprising/leading to errors, or silently do something different than what most people would expect. It requires A100 to have autocast in bf16.

This is due to not passing the `dtype` to `autocast ` here: https://github.com/pytorch/pytorch/blob/e2a4dfa468330c0587849bea4896ff5fffb33010/torch/cuda/amp/autocast_mode.py#L113-L124

Repro:
```python
import torch

class SomeFunc(torch.autograd.Function):
    @classmethod
    @torch.cuda.amp.custom_fwd
    def forward(cls, ctx, x):
        out = x.transpose(0, 1) @ x
        print("Forward dtype:", out.dtype)
        ctx.save_for_backward(out, x)
        return out

    @classmethod
    @torch.cuda.amp.custom_bwd
    def backward(cls, ctx, grad):
        out, x = ctx.saved_tensors
        out_recalc = x.transpose(0, 1) @ x
        print("Backward dtype:", out_recalc.dtype)
        return None

x = torch.randn([10, 10], device="cuda", requires_grad=True)
with torch.autocast("cuda", dtype=torch.bfloat16):
    y = SomeFunc.apply(x)
y.backward(torch.zeros_like(y))
```

Output:
```bash
Forward dtype: torch.bfloat16
Backward dtype: torch.float16
```

### Versions

```
Collecting environment information...
PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: 10.0.0-4ubuntu1 
CMake version: version 3.22.4
Libc version: glibc-2.31

Python version: 3.10.6 | packaged by conda-forge | (main, Aug 22 2022, 20:35:26) [GCC 10.4.0] (64-bit runtime)
Python platform: Linux-5.13.0-1025-aws-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.6.112
CUDA_MODULE_LOADING set to: 
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 510.47.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.982
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.1
[pip3] torch==1.12.1
[pip3] torchaudio==0.12.1
[pip3] torchmetrics==0.9.3
[pip3] torchvision==0.13.1
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.6.0              hecad31d_10    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h8d4b97c_729    conda-forge
[conda] mkl-service               2.4.0           py310ha2c4b55_0    conda-forge
[conda] mkl_fft                   1.3.1           py310h2b4bcf5_1    conda-forge
[conda] mkl_random                1.2.2           py310h00e6091_0  
[conda] numpy                     1.23.1          py310h1794996_0  
[conda] numpy-base                1.23.1          py310hcba007f_0  
[conda] pytorch                   1.12.1          py3.10_cuda11.6_cudnn8.3.2_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.12.1              py310_cu116    pytorch
[conda] torchmetrics              0.9.3                    pypi_0    pypi
[conda] torchvision               0.13.1              py310_cu116    pytorch
```

cc @mcarilli @ptrblck @leslie-fang-intel @jgong5