torch.jit.annotate — PyTorch 2.7 documentation (original) (raw)

torch.jit.annotate(the_type, the_value)[source][source]¶

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.