torch.jit.annotate

torch.jit.annotate(the_type, the_value)

Use to give type of the_value in TorchScript compiler.

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 Attribute() 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.