Repository: pytorch
Issue: Saving a LightningModule torch.jit.ScriptModule is incompatible with torch.amp.autocast
Labels: oncall: jit,has workaround,module: amp (automated mixed precision)

### 🐛 Describe the bug

I am training a model using pytorch lightning using mixed precision.  I have a region of the network that requires the dynamic range of float32, so I am using the `with torch.amp.autocast(enabled=False, device_type="cuda"):` to disable casting to lower precision. At the end of training, I save the nn.Module to torchscript format using `torch.jit.script`. However, I then encounter an error when running the scripted module.  This issue goes away when I remove the context manager.  The issue also doesn't occur using tracing instead of scripting.  However, we need to use scripting due to latency constraints. Please see full stack trace, and minimal example to reproduce below.

```python
import pytorch_lightning as pl
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchdata.datapipes.iter import IterableWrapper


class SimpleModel(nn.Module):
    def __init__(self, input_dim, hidden_dim=128, output_dim=2):
        super().__init__()
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.layer2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, inputs):
        with torch.amp.autocast(enabled=False, device_type="cuda"):
            x = self.layer1(inputs)
        x = self.layer2(x)
        return x


class SimpleTrainingModule(pl.LightningModule):
    def __init__(self, model):
         super().__init__()
         self.model = model
         self.loss = torch.nn.BCEWithLogitsLoss()

    def forward(self, inputs):
        return self.model(inputs)

    def training_step(self, inputs):
        input_tensor, target = inputs
        outputs = self.forward(input_tensor)
        loss = self.loss(outputs, target)
        return { "loss": loss }

    def configure_optimizers(self):
        return torch.optim.AdamW(self.model.parameters())


def data_generator(input_dim=4):
    for _ in range(20):
        yield torch.rand((1, input_dim), dtype=torch.float32), torch.empty((1, 2), dtype=torch.float32).random_(1)


class RandomDrivingDataModule(pl.LightningDataModule):
    def __init__(
        self,
        input_dim: int,
        batch_size: int,
    ) -> None:
        super().__init__()
        # no overhead in generating the datapipe here
        self.input_dim = input_dim
        self.batch_size = batch_size
        self.dp = IterableWrapper(data_generator(self.input_dim))

    def train_dataloader(self) -> DataLoader:
        return DataLoader(self.dp, batch_size=None, num_workers=1, pin_memory=True)

    def val_dataloader(self) -> DataLoader:
        return DataLoader(self.dp, batch_size=None, num_workers=1, pin_memory=True)


if __name__ == '__main__':
    trainer = pl.Trainer(
        accelerator="gpu",
        max_steps=10,
        devices=1,
    )
    datamodule = RandomDrivingDataModule(input_dim=4, batch_size=1)
    training_module = SimpleTrainingModule(SimpleModel(input_dim=4))
    trainer.fit(training_module, datamodule)

    # trace model
    example_input, _ = next(data_generator())
    scripted_module: torch.jit.ScriptModule = torch.jit.script(training_module)
    # scripted_module: torch.jit.ScriptModule = torch.jit.trace(training_module, example_inputs=example_input) <- This works!
    scripted_module(example_input)
```

The traceback I get is:
```python
Traceback (most recent call last):
  File "exception_example.py", line 83, in <module>
    saved_model(example_input)
  File "/home/sasha/si-venv-v2/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
RuntimeError: isInt() INTERNAL ASSERT FAILED at "../aten/src/ATen/core/ivalue.h":626, please report a bug to PyTorch. 
```

### Versions

Versions of relevant libraries:
[pip3] mypy-protobuf==3.3.0
[pip3] numpy==1.22.4
[pip3] pytorch-lightning==2.0.2
[pip3] torch==2.0.1
[pip3] torchdata==0.6.1
[pip3] torchmetrics==0.10.0
[pip3] torchvision==0.15.2
[pip3] triton==2.0.0
[conda] Could not collect

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel @mcarilli @ptrblck @leslie-fang-intel