torch.jit.annotate¶
-
torch.jit.
annotate
(the_type, the_value)[source]¶ This method is a pass-through function that returns the_value, used to hint TorchScript compiler the type of the_value. It is a no-op when running outside of TorchScript.
Though TorchScript can infer correct type for most Python expressions, there are some cases where type inference can be wrong, including:
Empty containers like [] and {}, which TorchScript assumes to be container of Tensor
Optional types like Optional[T] but assigned a valid value of type T, TorchScript would assume it is type T rather than Optional[T]
Note that annotate() does not help in __init__ method of torch.nn.Module subclasses because it is executed in eager mode. To annotate types of torch.nn.Module attributes, use
Annotate()
instead.Example:
import torch from typing import Dict @torch.jit.script def fn(): # Telling TorchScript that this empty dictionary is a (str -> int) dictionary # instead of default dictionary type of (str -> Tensor). d = torch.jit.annotate(Dict[str, int], {}) # Without `torch.jit.annotate` above, following statement would fail because of # type mismatch. d["name"] = 20
- Parameters
the_type – Python type that should be passed to TorchScript compiler as type hint for the_value
the_value – Value or expression to hint type for.
- Returns
the_value is passed back as return value.