JIT Function Fails when run a second time
Labels: oncall: jit

### 🐛 Describe the bug

JITed function fails after being run a second time on the same exact inputs with error: 
RuntimeError: tensor_type->scalarType().has_value() INTERNAL ASSERT FAILED. Missing Scalar Type information.

JIT Function Code:
```
import torch
from typing import List

@torch.jit.script
def sh_basis(degs: List[int], dirs, kappa=None):
    # evaluate a list of degrees
    if kappa is not None:
        kappa = kappa.reshape(-1)

    x, y, z = dirs.T[0], dirs.T[1], dirs.T[2]
    xx, yy, zz = x * x, y * y, z * z
    x4, y4, z4 = x**4, y**4, z**4
    x6, y6, z6 = x**6, y**6, z**6
    x8, y8, z8 = x**8, y**8, z**8
    x10, y10, z10 = x**10, y**10, z**10
    x12, y12, z12 = x**12, y**12, z**12
    x14, y14, z14 = x**14, y**14, z**14
    x16, y16, z16 = x**16, y**16, z**16

    values = []
    for deg in degs:
        scale = Al(deg, kappa) if kappa is not None else 1
        if deg == 0:
            values.append(scale*0.28209479177387814 * torch.ones_like(x))

        if deg == 1:
            values.extend([
                # scale*-0.690988*y,
                # scale*0.488603*z,
                # scale*-0.345494*x
                -scale*0.488603*x,
                scale*0.488603*z,
                -scale*0.488603*y
            ])

        if deg == 2:
            values.extend([
                scale*1.092548*y*x,
                -scale*1.092548*y*z,
                scale*0.315392*(3*zz - 1),
                -scale*1.092548*x*y,
                scale*0.546274*(xx-yy),
                # scale*1.5451*x*y,
                # scale*-1.5451*y*z,
                # scale*(0.946176*zz - 0.315392),
                # scale*-0.772548*x*z,
                # scale*(0.386274*xx - 0.386274*yy)
            ])

        if deg == 4:
            values.extend([
                2.50334*x*y*(xx - yy),
                -1.77013*y*z*(-3*xx + yy),
                0.946175*x*y*(7*zz - 1),
                0.669047*y*z*(7*zz - 3),
                3.70251*z4 - 3.17358*zz + 0.317358,
                0.669047*x*z*(7*zz - 3),
                (0.473087*xx - 0.473087*yy)*(7*zz - 1),
                1.77013*x*z*(xx - 3*yy),
                0.625836*x4 - 3.755016*xx*yy + 0.625836*y4,
            ])

        if deg == 8:
            values.extend([
                5.83141*x*y*(x6 - 7*x4*yy + 7*xx*y4 - y6),
                -2.91571*y*z*(-7*x6 + 35*x4*yy - 21*xx*y4 + y6),
                1.06467*x*y*(15*zz - 1)*(3*x4 - 10*xx*yy + 3*y4),
                3.44991*y*z*(5*zz - 1)*(5*x4 - 10*xx*yy + y4),
                1.91367*x*y*(xx - yy)*(65*z4 - 26*zz + 1),
                -1.23527*y*z*(-3*xx + yy)*(39*z4 - 26*zz + 3),
                0.912305*x*y*(143*z6 - 143*z4 + 33*zz - 1),
                0.109041*y*z*(715*z6 - 1001*z4 + 385*zz - 35),
                58.47336495*z8 - 109.15028124*z6 + 62.9713161*z4 - 11.4493302*zz + 0.31803695,
                0.109041*x*z*(715*z6 - 1001*z4 + 385*zz - 35),
                (0.456152*xx - 0.456152*yy)*(143*z6 - 143*z4 + 33*zz - 1),
                1.23527*x*z*(xx - 3*yy)*(39*z4 - 26*zz + 3),
                (0.478417*x4 - 2.870502*xx*yy + 0.478417*y4)*(65*z4 - 26*zz + 1),
                3.44991*x*z*(5*zz - 1)*(x4 - 10*xx*yy + 5*y4),
                (15*zz - 1)*(0.532333*x6 - 7.984995*x4*yy + 7.984995*xx*y4 - 0.532333*y6),
                2.91571*x*z*(x6 - 21*x4*yy + 35*xx*y4 - 7*y6),
                0.728927*x8 - 20.409956*x6*yy + 51.02489*x4*y4 - 20.409956*xx*y6 + 0.728927*y8,
            ])

        if deg == 16:
            values.extend([
                13.7174*x*y*(x14 - 35*x12*yy + 273*x10*y4 - 715*x8*y6 + 715*x6*y8 - 273*x4*y10 + 35*xx*y12 - y14),
                -4.84985*y*z*(-15*x14 + 455*x12*yy - 3003*x10*y4 + 6435*x8*y6 - 5005*x6*y8 + 1365*x4*y10 - 105*xx*y12 + y14),
                1.23186*x*y*(31*zz - 1)*(7*x12 - 182*x10*yy + 1001*x8*y4 - 1716*x6*y6 + 1001*x4*y8 - 182*xx*y10 + 7*y12),
                1.94775*y*z*(31*zz - 3)*(13*x12 - 286*x10*yy + 1287*x8*y4 - 1716*x6*y6 + 715*x4*y8 - 78*xx*y10 + y12),
                0.723375*x*y*(899*z4 - 174*zz + 3)*(3*x10 - 55*x8*yy + 198*x6*y4 - 198*x4*y6 + 55*xx*y8 - 3*y10),
                -0.427954*y*z*(899*z4 - 290*zz + 15)*(-11*x10 + 165*x8*yy - 462*x6*y4 + 330*x4*y6 - 55*xx*y8 + y10),
                0.20174*x*y*(8091*z6 - 3915*z4 + 405*zz - 5)*(5*x8 - 60*x6*yy + 126*x4*y4 - 60*xx*y6 + 5*y8),
                0.194401*y*z*(8091*z6 - 5481*z4 + 945*zz - 35)*(9*x8 - 84*x6*yy + 126*x4*y4 - 36*xx*y6 + y8),
                0.54985*x*y*(x6 - 7*x4*yy + 7*xx*y4 - y6)*(40455*z8 - 36540*z6 + 9450*z4 - 700*zz + 7),
                -0.336713*y*z*(-7*x6 + 35*x4*yy - 21*xx*y4 + y6)*(13485*z8 - 15660*z6 + 5670*z4 - 700*zz + 21),
                0.0444044*x*y*(3*x4 - 10*xx*yy + 3*y4)*(310155*z10 - 450225*z8 + 217350*z6 - 40250*z4 + 2415*zz - 21),
                0.0313986*y*z*(5*x4 - 10*xx*yy + y4)*(310155*z10 - 550275*z8 + 341550*z6 - 88550*z4 + 8855*zz - 231),
                0.166146*x*y*(xx - yy)*(310155*z12 - 660330*z10 + 512325*z8 - 177100*z6 + 26565*z4 - 1386*zz + 11),
                -0.0515197*y*z*(-3*xx + yy)*(310155*z12 - 780390*z10 + 740025*z8 - 328900*z6 + 69069*z4 - 6006*zz + 143),
                0.00631775*x*y*(5892950.0*z14 - 17298600.0*z12 + 19684700.0*z10 - 10935900.0*z8 + 3062060.0*z6 - 399399*z4 + 19019*zz - 143),
                0.00115346*y*z*(17678800.0*z14 - 59879900.0*z12 + 80528200.0*z10 - 54679600.0*z8 + 19684700.0*z6 - 3594590.0*z4 + 285285*zz - 6435),
                14862.935214*z16 - 57533.910858*z14 + 90269.063271*z12 - 73552.588389*z10 + 33098.5905939*z8 - 8058.7928655*z6 + 959.37986754*z4 - 43.280250156*zz + 0.3182371335,
                0.00115346*x*z*(17678800.0*z14 - 59879900.0*z12 + 80528200.0*z10 - 54679600.0*z8 + 19684700.0*z6 - 3594590.0*z4 + 285285*zz - 6435),
                (0.00315887*xx - 0.00315887*yy)*(5892950.0*z14 - 17298600.0*z12 + 19684700.0*z10 - 10935900.0*z8 + 3062060.0*z6 - 399399*z4 + 19019*zz - 143),
                0.0515197*x*z*(xx - 3*yy)*(310155*z12 - 780390*z10 + 740025*z8 - 328900*z6 + 69069*z4 - 6006*zz + 143),
                (0.0415365*x4 - 0.249219*xx*yy + 0.0415365*y4)*(310155*z12 - 660330*z10 + 512325*z8 - 177100*z6 + 26565*z4 - 1386*zz + 11),
                0.0313986*x*z*(x4 - 10*xx*yy + 5*y4)*(310155*z10 - 550275*z8 + 341550*z6 - 88550*z4 + 8855*zz - 231),
                (0.0222022*x6 - 0.333033*x4*yy + 0.333033*xx*y4 - 0.0222022*y6)*(310155*z10 - 450225*z8 + 217350*z6 - 40250*z4 + 2415*zz - 21),
                0.336713*x*z*(x6 - 21*x4*yy + 35*xx*y4 - 7*y6)*(13485*z8 - 15660*z6 + 5670*z4 - 700*zz + 21),
                (0.0687312*x8 - 1.9244736*x6*yy + 4.811184*x4*y4 - 1.9244736*xx*y6 + 0.0687312*y8)*(40455*z8 - 36540*z6 + 9450*z4 - 700*zz + 7),
                0.194401*x*z*(8091*z6 - 5481*z4 + 945*zz - 35)*(x8 - 36*x6*yy + 126*x4*y4 - 84*xx*y6 + 9*y8),
                (8091*z6 - 3915*z4 + 405*zz - 5)*(0.10087*x10 - 4.53915*x8*yy + 21.1827*x6*y4 - 21.1827*x4*y6 + 4.53915*xx*y8 - 0.10087*y10),
                0.427954*x*z*(899*z4 - 290*zz + 15)*(x10 - 55*x8*yy + 330*x6*y4 - 462*x4*y6 + 165*xx*y8 - 11*y10),
                (899*z4 - 174*zz + 3)*(0.180844*x12 - 11.935704*x10*yy + 89.51778*x8*y4 - 167.099856*x6*y6 + 89.51778*x4*y8 - 11.935704*xx*y10 + 0.180844*y12),
                1.94775*x*z*(31*zz - 3)*(x12 - 78*x10*yy + 715*x8*y4 - 1716*x6*y6 + 1287*x4*y8 - 286*xx*y10 + 13*y12),
                (31*zz - 1)*(0.615932*x14 - 56.049812*x12*yy + 616.547932*x10*y4 - 1849.643796*x8*y6 + 1849.643796*x6*y8 - 616.547932*x4*y10 + 56.049812*xx*y12 - 0.615932*y14),
                4.84985*x*z*(x14 - 105*x12*yy + 1365*x10*y4 - 5005*x8*y6 + 6435*x6*y8 - 3003*x4*y10 + 455*xx*y12 - 15*y14),
                0.857341*x16 - 102.88092*x14*yy + 1560.36062*x12*y4 - 6865.586728*x10*y6 + 11033.97867*x8*y8 - 6865.586728*x6*y10 + 1560.36062*x4*y12 - 102.88092*xx*y14 + 0.857341*y16,
            ])
    return torch.stack(values, dim=-1)

```

How to reproduce the error:
```
Python 3.10.5 | packaged by conda-forge | (main, Jun 14 2022, 07:06:46) [GCC 10.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> from modules.sh import sh_basis
>>> device = torch.device('cuda')
>>> vecs = torch.rand((10000, 3), device=device)
>>> kappa = torch.rand((10000, 1), device=device)
>>> sh_basis([0, 1, 2, 4], vecs, kappa)
tensor([[ 2.8209e-01, -7.4768e-02,  1.0519e-01,  ..., -1.8813e-03,
         -8.9227e-01, -4.4803e-01],
        [ 2.8209e-01, -4.2040e-05,  1.3623e-04,  ..., -2.3001e-01,
         -2.3584e-01,  1.1216e-01],
        [ 2.8209e-01, -8.7236e-03,  1.7843e-02,  ..., -2.0500e-02,
         -6.4801e-02, -1.3913e-02],
        ...,
        [ 2.8209e-01, -2.6438e-02,  9.2806e-02,  ...,  1.4510e-01,
          1.8478e-02,  7.0536e-04],
        [ 2.8209e-01, -7.5585e-12,  7.9083e-12,  ...,  5.1023e-01,
         -8.5628e-01, -7.5704e-01],
        [ 2.8209e-01, -2.9055e-02,  1.1445e-01,  ...,  3.6547e-02,
         -1.2329e-03, -1.0240e-03]], device='cuda:0')
>>> sh_basis([0, 1, 2, 4], vecs, kappa)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: tensor_type->scalarType().has_value() INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1659484803030/work/torch/csrc/jit/codegen/cuda/type_promotion.cpp":111, please report a bug to PyTorch. Missing Scalar Type information
```


### Versions

PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Manjaro Linux (x86_64)
GCC version: (GCC) 12.2.0
Clang version: 14.0.6
CMake version: version 3.24.0
Libc version: glibc-2.36

Python version: 3.10.5 | packaged by conda-forge | (main, Jun 14 2022, 07:06:46) [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.15.85-1-MANJARO-x86_64-with-glibc2.36
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to:
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090
Nvidia driver version: 525.60.11
cuDNN version: Probably one of the following:
/usr/lib/libcudnn.so.8.6.0
/usr/lib/libcudnn_adv_infer.so.8.6.0
/usr/lib/libcudnn_adv_train.so.8.6.0
/usr/lib/libcudnn_cnn_infer.so.8.6.0
/usr/lib/libcudnn_cnn_train.so.8.6.0
/usr/lib/libcudnn_ops_infer.so.8.6.0
/usr/lib/libcudnn_ops_train.so.8.6.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.1
[pip3] pytorch-lightning==1.7.7
[pip3] torch==1.12.1
[pip3] torch-ema==0.3
[pip3] torch-scatter==2.0.9
[pip3] torchaudio==0.12.1
[pip3] torchmetrics==0.10.2
[pip3] torchvision==0.13.1
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.6.0              hecad31d_10    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libblas                   3.9.0            12_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            12_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            12_linux64_mkl    conda-forge
[conda] mkl                       2021.4.0           h06a4308_640
[conda] mkl-service               2.4.0           py310h7f8727e_0
[conda] mkl_fft                   1.3.1           py310hd6ae3a3_0
[conda] mkl_random                1.2.2           py310h00e6091_0
[conda] numpy                     1.22.4                   pypi_0    pypi
[conda] numpy-base                1.23.1          py310hcba007f_0
[conda] pytorch                   1.12.1          py3.10_cuda11.6_cudnn8.3.2_0    pytorch
[conda] pytorch-lightning         1.7.7                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch-scatter           2.0.9           py310_torch_1.12.0_cu116    pyg
[conda] torch-ema                 0.3                      pypi_0    pypi
[conda] torchaudio                0.12.1              py310_cu116    pytorch
[conda] torchmetrics              0.10.2                   pypi_0    pypi
[conda] torchvision               0.13.1              py310_cu116    pytorch

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel