torch._dynamo.exc.TorchRunTimeError in get_fake_value while performing quantization aware training
Labels: oncall: quantization, triaged, oncall: pt2

### 🐛 Describe the bug

I've been trying to implement quantization-aware training on a convnextSmall model. It errors out saying RuntimeError: Failed running call_module self_loss_fn(*(FakeTensor(FakeTensor(..., device='meta', size=(256, 1)), cuda:0), FakeTensor(FakeTensor(..., device='meta', size=(256,), dtype=torch.float64), cuda:0)), **{}):    RuntimeError: gather(): Expected dtype int64 for index, but got torch.float64 . I'm not sure what is causing this error or how to fix it. Any leads about what is going on here will be very helpful.

**Sample code to reproduce the problem**

```
class ConvNextSmallQAT(pl.LightningModule):
    def __init__(self, num_classes):
        super().__init__()
        self.model = convnext_small(weights="DEFAULT")
        self.model.classifier[2] = torch.nn.Linear(
            in_features=self.model.classifier[2].in_features, out_features=num_classes
        )
        self.quant = QuantStub()
        self.dequant = DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.model(x)
        x = self.dequant(x)
        return x

class ExperimentExectutor:
    def __init__(
        self,
        experiment_config,
        num_gpus: int = 1,
        num_workers: int = 8,
        convert_model: bool = True,
    ):
        self.batch_size: int = experiment_config.batch_size
        self.max_steps: int = experiment_config.max_steps
        self.optimizer: torch.optim = experiment_config.optimizer
        self.model: torch.nn.Module = experiment_config.model
        self.model_settings: dict = experiment_config.model_settings
        self.data: pl.LightningDataModule = experiment_config.data
        self.num_gpus: int = num_gpus
        self.num_workers: int = num_workers
        self.convert_model: bool = convert_model
        self.experiment_config_name = experiment_config().__class__.__name__

    def __call__(self):
        dm = self.data(
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            num_gpus=self.num_gpus,
        )
        if self.prepare_data:
            dm.prepare_data()
            dm.setup()

        model = self.trainer(
            self.model,
            self.model_settings,
            self.label_smoothing,
            self.optimizer_settings,
            self.optimizer,
            self.lr_scheduler_settings,
            self.lr_scheduler,
        )
        compiled_model = torch.compile(model)
        tensorboard_logger = TensorBoardLogger("../outputs/tensorboard/")
        trainer = pl.Trainer(
            max_steps=self.max_steps,
            val_check_interval=self.validation_frequency,
            logger=[ tensorboard_logger],
            accelerator="auto",
            callbacks=[
                checkpoint_callback,
                TQDMProgressBar(refresh_rate=10),
                LearningRateMonitor(logging_interval="step"),
            ],
            devices=self.num_gpus,
            # precision=16,
        )
        print("Before QAT:")
        print_size_of_model(compiled_model)
        compiled_model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
        torch.quantization.prepare_qat(compiled_model, inplace=True)

        trainer.fit(compiled_model, dm)
```

**Error message**
```
 WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
 Downloading: "https://download.pytorch.org/models/convnext_small-0c510722.pth" to /root/.cache/torch/hub/checkpoints/convnext_small-0c510722.pth
100%|██████████| 192M/192M [00:04<00:00, 47.9MB/s] 
 GPU available: True (cuda), used: True
 TPU available: False, using: 0 TPU cores
 IPU available: False, using: 0 IPUs
 HPU available: False, using: 0 HPUs
Before QAT:
 model size: 188.658MB
 /usr/local/lib/python3.10/dist-packages/torch/ao/quantization/observer.py:214: UserWarning: Please use quant_min and quant_max to specify the range for observers.                     reduce_range will be deprecated in a future release of PyTorch.
   warnings.warn(
 LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
   | Name                | Type             | Params
 ---------------------------------------------------------

Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/utils.py", line 1199, in run_node
    return nnmodule(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1501, in _call_impl
     return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/loss.py", line 1174, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py", line 3029, in cross_entropy
     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
   File "/usr/local/lib/python3.10/dist-packages/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/torch/_subclasses/fake_tensor.py", line 987, in __torch_dispatch__
     return self.dispatch(func, types, args, kwargs)
   File "/usr/local/lib/python3.10/dist-packages/torch/_subclasses/fake_tensor.py", line 1140, in dispatch
    return decomposition_table[func](*args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/torch/_decomp/decompositions.py", line 2804, in nll_loss_forward
     result = -torch.gather(self, channel_dim, safe_target_).squeeze(channel_dim)
   File "/usr/local/lib/python3.10/dist-packages/torch/utils/_stats.py", line 20, in wrapper
     return fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/_subclasses/fake_tensor.py", line 987, in __torch_dispatch__
     return self.dispatch(func, types, args, kwargs)
   File "/usr/local/lib/python3.10/dist-packages/torch/_subclasses/fake_tensor.py", line 1170, in dispatch
     r = func(*args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 287, in __call__
     return self._op(*args, **kwargs or {})
  File "/usr/local/lib/python3.10/dist-packages/torch/_meta_registrations.py", line 1868, in meta_gather
    check(
   File "/usr/local/lib/python3.10/dist-packages/torch/_prims_common/__init__.py", line 1563, in check
     raise exc_type(s())
 RuntimeError: gather(): Expected dtype int64 for index, but got torch.float64

The above exception was the direct cause of the following exception:

 Traceback (most recent call last):
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/utils.py", line 1152, in get_fake_value
    return wrap_fake_exception(
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/utils.py", line 808, in wrap_fake_exception
    return fn()
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/utils.py", line 1153, in <lambda>
    lambda: run_node(tx.output, node, args, kwargs, nnmodule)
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/utils.py", line 1206, in run_node
     raise RuntimeError(
 RuntimeError: Failed running call_module self_loss_fn(*(FakeTensor(FakeTensor(..., device='meta', size=(256, 1)), cuda:0), FakeTensor(FakeTensor(..., device='meta', size=(256,), dtype=torch.float64), cuda:0)), **{}):
 gather(): Expected dtype int64 for index, but got torch.float64
(scroll up for backtrace)
 
 The above exception was the direct cause of the following exception:
Traceback` (most recent call last):
   File "/plx-context/artifacts/ee60add6a77f4cf289ffde3ab9b75f41/uploads/main.py", line 35, in <module>  
        main()
   File "/plx-context/artifacts/ee60add6a77f4cf289ffde3ab9b75f41/uploads/main.py", line 29, in main   
         ExperimentExectutor(
  File "/plx-context/artifacts/ee60add6a77f4cf289ffde3ab9b75f41/uploads/experiment_executor_qat.py", line 106, in __call__
        trainer.fit(compiled_model, dm)
   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 531, in fit
    call._call_and_handle_interrupt(
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 570, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 975, in _run
    results = self._run_stage()
   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1016, in _run_stage
     self._run_sanity_check()
   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1045, in _run_sanity_check
     val_loop.run()
   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/utilities.py", line 177, in _decorator
     return loop_run(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 115, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx)
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 375, in _evaluation_step
     output = call._call_strategy_hook(trainer, hook_name, *step_kwargs.values())
   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 287, in _call_strategy_hook
     output = fn(*args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/strategy.py", line 379, in validation_step
     return self.model.validation_step(*args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
     return fn(*args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py", line 337, in catch_errors
     return callback(frame, cache_size, hooks)
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py", line 404, in _convert_frame
     result = inner_convert(frame, cache_size, hooks)
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py", line 104, in _fn
     return fn(*args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py", line 262, in _convert_frame_assert
    return _compile(
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/utils.py", line 163, in time_wrapper
     r = func(*args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py", line 324, in _compile
     out_code = transform_code_object(code, transform)
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/bytecode_transformation.py", line 445, in transform_code_object
     transformations(instructions, code_options)
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py", line 311, in transform
     tracer.run()
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py", line 1726, in run
     super().run()
File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py", line 576, in run.   
     and self.step().   
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py", line 540, in step
    getattr(self, inst.opname)(inst)
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py", line 342, in wrapper
     return inner_fn(self, inst)
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py", line 965, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py", line 474, in call_function
     self.push(fn.call_function(self, args, kwargs))
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/variables/nn_module.py", line 203, in call_function
     return wrap_fx_proxy(
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/variables/builder.py", line 754, in wrap_fx_proxy
     return wrap_fx_proxy_cls(
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/variables/builder.py", line 789, in wrap_fx_proxy_cls
     example_value = get_fake_value(proxy.node, tx)
   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/utils.py", line 1173, in get_fake_value
    raise TorchRuntimeError() from e
 torch._dynamo.exc.TorchRuntimeError:

 from user code:
    File "/plx-context/artifacts/ee60add6a77f4cf289ffde3ab9b75f41/uploads/trainers/classification.py", line 194, in validation_step
     loss = self.loss_fn(logits, y)

 Set torch._dynamo.config.verbose=True for more information
 You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True``` 

### Versions

pytorch-lightning-2.0.3
torchvision==0.15.2
requests-2.31.0 
setuptools-59.5.0 
sympy-1.12 
torch-2.0.1+cu117 
torchmetrics-0.10.3 
tqdm-4.65.0 
triton-2.0.0 
typing-extensions-4.4.0 
urllib3-1.26.16 
```

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel @ezyang @msaroufim @wconstab @bdhirsh @anijain2305