dict.py

from torch import nn, jit
from typing import Dict

@jit.script
def lookup_all(d, nested_keys):
    # type: (Dict[str, Tensor], List[str]) -> List[Tensor]
    tensors = []
    for i in range(len(nested_keys)):
        key = nested_keys[i]
        tensors.append(d[key])
    return tensors

class MyDictModel(jit.ScriptModule):
    @jit.script_method
    def forward(self, keys, lookup):
        # type: (List[str], Dict[str, Tensor]) -> List[Tensor]
        return lookup_all(lookup, keys)


vocab = {chr(o): torch.LongTensor(o) for o in range(ord('a'), ord('z') + 1)}
model = MyDictModel()
input = list('cheese')
print(model(input, vocab))

model.save("my_model.pt")