Repository: pytorch
Issue: [LTC] OOM on mnist example
Labels: triaged,module: lazy

### 🐛 Describe the bug

OOM on the example mnist.py while enabling CUDA device, the error logs is as the following. I just use `torch.cuda.memory_stats` to print the allocated CUDA memory, and we can see the memory is increasing:


``` text
Train Epoch: 1 [0/60000 (0%)]	Loss: 2.305605
allocated_bytes.all.current: 33812992
Train Epoch: 1 [640/60000 (1%)]	Loss: 1.253851
allocated_bytes.all.current: 180380160
....
Train Epoch: 2 [0/60000 (0%)]	Loss: 0.034653
allocated_bytes.all.current: 13733059072
Train Epoch: 2 [640/60000 (1%)]	Loss: 0.074930
allocated_bytes.all.current: 13879101952
Train Epoch: 2 [1280/60000 (2%)]	Loss: 0.003717
allocated_bytes.all.current: 14025144832
Traceback (most recent call last):
  File "test_mnist.py", line 86, in <module>
    train(log_interval, model, device, train_loader, optimizer, epoch)
  File "test_mnist.py", line 49, in train
    torch._lazy.mark_step()
  File "/workspace/BladeDISC/venv-py38/lib/python3.8/site-packages/torch/_lazy/__init__.py", line 12, in mark_step
    torch._C._lazy._mark_step(device, [], wait=wait)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 14.76 GiB total capacity; 13.18 GiB already allocated; 9.75 MiB free; 13.66 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```

## Reproduce Steps
 
``` bash
git clone git@github.com:pytorch/pytorch.git
LTC_TS_CUDA=1 python pytorch/torch/csrc/lazy/test_mnist.py 
```

### Versions

Collecting environment information...
PyTorch version: 1.12.0+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
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: Could not collect
CMake version: version 3.20.0
Libc version: glibc-2.31

Python version: 3.8.10 (default, Mar 15 2022, 12:22:08)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-4.15.0-162-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.3.109
GPU models and configuration:
GPU 0: Tesla T4
GPU 1: Tesla T4

Nvidia driver version: 460.91.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.0
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_train.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_train.so.8
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.0
[pip3] torch==1.12.0+cu113
[conda] Could not collect


cc @wconstab