Repository: pytorch
Issue: dtype mismatch when after using auto mixed precision
Labels: triaged,module: amp (automated mixed precision)

### 🐛 Describe the bug

I am training model that includes 2 fully connected networks.
So far I have not seen any issue with them, and I wanted to use a mixed precision training to accelerate the process.
I made the required changes to the code but strangely, after several epochs of successful training and loss curve behavious, I get the following error:
```
Traceback (most recent call last):
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2021.1.2\plugins\python-ce\helpers\pydev\pydevd.py", line 1483, in _exec
    pydev_imports.execfile(file, globals, locals)  # execute the script
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2021.1.2\plugins\python-ce\helpers\pydev\_pydev_imps\_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "<path to my code>.py", line 2146, in <module>
    train()
  File "<path to my code>.py", line 2007, in train
    scaler.scale(loss_total).backward()
  File "C:\Users\shama\Anaconda3\envs\pytorch\lib\site-packages\torch\_tensor.py", line 255, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "C:\Users\shama\Anaconda3\envs\pytorch\lib\site-packages\torch\autograd\__init__.py", line 147, in backward
    Variable._execution_engine.run_backward(
RuntimeError: masked_scatter: expected self and source to have same dtypes but gotHalf and Float
```
I wonder why would that happen in the middle of training ( after 20% of the training )

The code is huge, but here's its basic structure:
```
 model = Network()
 # different params trained with different learning rate, which is set in the training loop
 optimizer = pt.optim.Adam([
      {'params': paramsA, 'lr': 0},
      {'params': paramsB, 'lr': 0},
      {'params': paramsC, 'lr': 0} ])
 writer = SummaryWriter(..)
 for epoch in range(start_epoch, args.epochs+1):
    epoch_loss_total = 0
    epoch_mse = 0

    model.train()

    for i, sample in enumerate(dataloader_train):
      # print("step: {}".format(i))
      setLearningRate(optimizer, epoch)
      optimizer.zero_grad()
      
      sel = selectPixels()   

      gt = sample ['image']
      gt = gt.view(gt.shape[0], gt.shape[1], gt.shape[2] * gt.shape[3])
      gt = gt[:, :, sel, None].cuda()

      with autocast(enabled=args.mixed_precision):

        output, in_range_selection = model(sample, sel)
        gt = gt[:, :, in_range_selection,:]

        mse = pt.mean((output - gt) ** 2)

        loss_total = mse

        loss2 = 0
        if SomeCondition1():
          loss2= args.l2_weight * pt.mean((pt.sigmoid(model.feature)))        

        loss_total = loss_total + loss2

        # another loss
        loss_total = loss_total + args.l3_weight * func(gt, output)

        # for loss statistics
        epoch_loss_total += loss_total
        epoch_mse += mse
      
      # this section is out of autocast
      scaler.scale(loss_total).backward()

      scaler.step(optimizer)

      scaler.update()

      step += 1
    # this section is out of samples loop 
    if someCondition2():
      with pt.no_grad():
        out_image = evaluateResults1(model)

        writer.add_image('images/res', out_image , 2), epoch)        
        pt.cuda.empty_cache()

    if someCondition3():
      dumpResults(model)
      pt.cuda.empty_cache()

    # report mean value
    mean = pt.mean(model.mlp)
    writer.add_scalar('loss/mean', mean, epoch)

    if someCondition4():
      checkpoint(ckpt, model, optimizer, epoch+1)
```

### Versions

Collecting environment information...
PyTorch version: 1.9.0
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Home
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.8.10 (default, May 19 2021, 13:12:57) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19044-SP0
Is CUDA available: True
CUDA runtime version: 11.3.109
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3070
Nvidia driver version: 472.39
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.20.2
[pip3] torch==1.9.0
[pip3] torch-tb-profiler==0.2.1
[pip3] torchvision==0.10.0
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.1.1               heb2d755_7    conda-forge
[conda] mkl                       2021.2.0           haa95532_296
[conda] mkl-service               2.3.0            py38h2bbff1b_1
[conda] mkl_fft                   1.3.0            py38h277e83a_2
[conda] mkl_random                1.2.1            py38hf11a4ad_2
[conda] numpy                     1.20.2           py38ha4e8547_0
[conda] numpy-base                1.20.2           py38hc2deb75_0
[conda] pytorch                   1.9.0           py3.8_cuda11.1_cudnn8_0    pytorch
[conda] torch-tb-profiler         0.2.1                    pypi_0    pypi
[conda] torchvision               0.10.0               py38_cu111    pytorch


cc @mcarilli @ptrblck