Repository: pytorch
Issue: AMP consumes 30x gpu memory with bmm
Labels: triaged,module: amp (automated mixed precision)

### 🐛 Describe the bug

AMP consumes about 30x gpu memory when used on `bmm`.

Code:
```
import torch
import torch.nn as nn


class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        base = torch.randn((41292, 64))
        self.register_buffer("base", base)

    def forward(self, tensor):
        batch_size = tensor.shape[0]
        tensor = tensor.unsqueeze(-1)  # (B, 64, 1)
        return self.base.expand(batch_size, -1, -1).bmm(tensor)


device = torch.device('cuda')
model = MyModule().to(device)

for use_amp in [True, False]:
    for i in range(10):
        with torch.autocast(
            device_type="cuda", dtype=torch.float16, enabled=use_amp
        ):
            output = model(torch.randn((1024, 64), requires_grad=True, device=device))
            torch.cuda.empty_cache()
            print("GPU Memory Used:", torch.cuda.memory_allocated())
```

Result:
```
GPU Memory Used: 5508148224
GPU Memory Used: 5508148224
GPU Memory Used: 5508148224
GPU Memory Used: 5508148224
GPU Memory Used: 5508148224
GPU Memory Used: 5508148224
GPU Memory Used: 5508148224
GPU Memory Used: 5508148224
GPU Memory Used: 5508148224
GPU Memory Used: 5508148224
GPU Memory Used: 180702208
GPU Memory Used: 180702208
GPU Memory Used: 180702208
GPU Memory Used: 180702208
GPU Memory Used: 180702208
GPU Memory Used: 180702208
GPU Memory Used: 180702208
GPU Memory Used: 180702208
GPU Memory Used: 180702208
GPU Memory Used: 180702208
```

### Versions

```
Collecting environment information...
PyTorch version: 1.10.0+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.22.4
Libc version: glibc-2.10

Python version: 3.7.10 (default, Feb 26 2021, 18:47:35)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-126-generic-x86_64-with-debian-buster-sid
Is CUDA available: True
CUDA runtime version: 11.1.105
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090

Nvidia driver version: 510.85.02
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.6
[pip3] pytorch3d==0.6.2
[pip3] torch==1.10.0+cu113
[pip3] torch-tb-profiler==0.4.0
[pip3] torchaudio==0.10.0+cu113
[pip3] torchfile==0.1.0
[pip3] torchvision==0.11.1+cu113
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.1.74              h6bb024c_0    nvidia
[conda] cudatoolkit-dev           11.4.0               h5e8e339_5    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640
[conda] mkl-fft                   1.3.0                    pypi_0    pypi
[conda] mkl-random                1.2.1                    pypi_0    pypi
[conda] mkl-service               2.3.0                    pypi_0    pypi
[conda] mkl_fft                   1.3.1            py37hd3c417c_0
[conda] mkl_random                1.2.2            py37h51133e4_0
[conda] numpy                     1.20.2                   pypi_0    pypi
[conda] numpy-base                1.21.5           py37ha15fc14_3
[conda] pytorch-msssim            0.2.1                    pypi_0    pypi
[conda] torch                     1.10.0+cu113             pypi_0    pypi
[conda] torch-tb-profiler         0.4.0                    pypi_0    pypi
[conda] torchaudio                0.10.0+cu113             pypi_0    pypi
[conda] torchelastic              0.2.0                    pypi_0    pypi
[conda] torchfile                 0.1.0                    pypi_0    pypi
[conda] torchtext                 0.10.0                     py37    pytorch
[conda] torchvision               0.11.1+cu113             pypi_0    pypi
```

cc @mcarilli @ptrblck