How to specify key/value types for IValue::GenericDict?

[Contextualist]

I am trying to run a torchscript model with .method_is. The model has a method defined as follows:

    @torch.jit.export
    def generate(self, batch: Dict[str, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
        ...

I am trying to construct an input as follows:

    let input = IValue::GenericDict(vec![                                                             
        (IValue::String("text".to_owned()), IValue::Tensor(Tensor::of_slice(&[/* ... */]).reshape(&[-1, 2]))),  
        (IValue::String("text_len".to_owned()), IValue::Tensor(Tensor::of_slice(&[29, 38]))),
        (IValue::String("start_index".to_owned()), IValue::Tensor(Tensor::of_slice(&[1, 1]))),  
    ]);
    let output = model.method_is("generate", &[input])?;

The attempt above failed and produced a runtime error:

Error: Internal torch error: generate() Expected a value of type 'Dict[str, Tensor]' for argument 'batch' but instead found type 'Dict[Any, Any]'.
Position: 1
Declaration: generate(__torch__.dp.model.model.ForwardTransformer self, Dict(str, Tensor) batch) -> ((Tensor, Tensor))
Exception raised from checkArg at /path/to/libtorch/include/ATen/core/function_schema_inl.h:336 (most recent call first):

Is it possible to specify key/value types for GenericDict, so that it won't be recognized as Dict[Any, Any]?

[LaurentMazare]

There wasn't a good way to do this in the current version so I've made some changes in #598 that should specialize the dictionaries in this case. That said, I'm not very happy with the implementation and it's specific to this type of dictionaries so I plan on refactoring this a bit but hopefully this should already cover your use case (which seems to be the most common use case for such dictionaries).
You can see this being used in a newly added example here.

[Contextualist]

Thanks for your prompt action! And it does cover my use case. Feel free to close this issue or leave it open as a tracking issue.

[LaurentMazare]

Closing this for now, feel free to re-open if further issues.