Repository: pytorch
Issue: Torch.jit : RuntimeError:  Unable to extract string literal index for ModuleDict
Labels: oncall: jit

### 🐛 Describe the bug

Hi, I'm trying to use torch-mlir to compile a pytorch wide_and_deep model from https://github.com/jrzaurin/pytorch-widedeep. But I'm having some issues now, when compiling the model. According to reply from torch-mlir, the error due to torch.jit.script that can not extract string literal index for ModuleDict .  When I call torch_mlir.compile, torch-mlir calls torch.jit.script on the module before importing it into torch-mlir.

1. Python Script
```
import pandas as pd
import numpy as np
import torch
from sklearn.model_selection import train_test_split

from pytorch_widedeep import Trainer
from pytorch_widedeep.preprocessing import WidePreprocessor, TabPreprocessor
from pytorch_widedeep.models import Wide, TabMlp, WideDeep
from pytorch_widedeep.metrics import Accuracy
from pytorch_widedeep.datasets import load_adult

import torch_mlir

df = load_adult(as_frame=True)
df["income_label"] = (df["income"].apply(lambda x: ">50K" in x)).astype(int)
df.drop("income", axis=1, inplace=True)
df_train, df_test = train_test_split(df, test_size=0.2, stratify=df.income_label)

# Define the 'column set up'
wide_cols = [
    "education",
    "relationship",
    "workclass",
    "occupation",
    "native-country",
    "gender",
]
crossed_cols = [("education", "occupation"), ("native-country", "occupation")]

cat_embed_cols = [
    "workclass",
    "education",
    "marital-status",
    "occupation",
    "relationship",
    "race",
    "gender",
    "capital-gain",
    "capital-loss",
    "native-country",
]
continuous_cols = ["age", "hours-per-week"]
target = "income_label"
target = df_train[target].values

# prepare the data
wide_preprocessor = WidePreprocessor(wide_cols=wide_cols, crossed_cols=crossed_cols)
X_wide = wide_preprocessor.fit_transform(df_train)

tab_preprocessor = TabPreprocessor(
    cat_embed_cols=cat_embed_cols, continuous_cols=continuous_cols  # type: ignore[arg-type]
)
X_tab = tab_preprocessor.fit_transform(df_train)

# build the model
wide = Wide(input_dim=np.unique(X_wide).shape[0], pred_dim=1)
tab_mlp = TabMlp(
    column_idx=tab_preprocessor.column_idx,
    cat_embed_input=tab_preprocessor.cat_embed_input,
    continuous_cols=continuous_cols,
)
model = WideDeep(wide=wide, deeptabular=tab_mlp)
print("dim: ",np.unique(X_wide).shape[0])
# train and validate
trainer = Trainer(model, objective="binary", metrics=[Accuracy])
trainer.fit(
    X_wide=X_wide,
    X_tab=X_tab,
    target=target,
    n_epochs=5,
    batch_size=256,
)

# predict on test
X_wide_te = wide_preprocessor.transform(df_test)
X_tab_te = tab_preprocessor.transform(df_test)
preds = trainer.predict(X_wide=X_wide_te, X_tab=X_tab_te)

# Save and load

# Option 1: this will also save training history and lr history if the
# LRHistory callback is used
trainer.save(path="model_weights", save_state_dict=True)

# Option 2: save as any other torch model
torch.save(model.state_dict(), "model_weights/wd_model.pt")

# From here in advance, Option 1 or 2 are the same. I assume the user has
# prepared the data and defined the new model components:
# 1. Build the model
model_new = WideDeep(wide=wide, deeptabular=tab_mlp)
model_new.load_state_dict(torch.load("model_weights/wd_model.pt"))

# 2. Instantiate the trainer
trainer_new = Trainer(model_new, objective="binary")

# 3. Either start the fit or directly predict
preds = trainer_new.predict(X_wide=X_wide, X_tab=X_tab)

data = []
data.append(torch.from_numpy(X_wide))
data.append(torch.from_numpy(X_tab))
print(model_new)

module = torch_mlir.compile(model_new, data, output_type=torch_mlir.OutputType.LINALG_ON_TENSORS, use_tracing=False)
```

2. Error
```
Traceback (most recent call last):
  File "/home/project/iree-demo/pytorch/test-case-generate/widedeep/wide_deep.py", line 105, in <module>
    module = torch_mlir.compile(model_new, data, output_type=torch_mlir.OutputType.LINALG_ON_TENSORS, use_tracing=False)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch_mlir/__init__.py", line 404, in compile
    scripted = torch.jit.script(model)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_script.py", line 1284, in script
    return torch.jit._recursive.create_script_module(
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 480, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_script.py", line 614, in _construct
    init_fn(script_module)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 520, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_script.py", line 614, in _construct
    init_fn(script_module)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 520, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_script.py", line 614, in _construct
    init_fn(script_module)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 520, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_script.py", line 614, in _construct
    init_fn(script_module)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 520, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 546, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/torch/jit/_recursive.py", line 397, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError: 
Unable to extract string literal index. ModuleDict indexing is only supported with string literals. For example, 'i = "a"; self.layers[i](x)' will fail because i is not a literal. Enumeration of ModuleDict is supported, e.g. 'for k, v in self.items(): out = v(inp)':
  File "/root/anaconda3/envs/torch-mlir/lib/python3.10/site-packages/pytorch_widedeep/models/tabular/embeddings_layers.py", line 173
    def forward(self, X: Tensor) -> Tensor:
        embed = [
            self.embed_layers["emb_layer_" + self.embed_layers_names[col]](
            ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
                X[:, self.column_idx[col]].long()
            )

```

### Versions

```
conda create -n torch-mlir python=3.10
pip install torch-2.1.0.dev20230523+cpu-cp310-cp310-linux_x86_64.whl
pip install torch_mlir-20230524.848-cp310-cp310-linux_x86_64.whl
pip install torchvision-0.16.0.dev20230523+cpu-cp310-cp310-linux_x86_64.whl
pip install pytorch-widedeep
```

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel