Rate this Page
★★★★★

torch.jit.annotate#

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

Use to give type ofthe_value in TorchScript compiler.

This method is a pass-through function that returnsthe_value, used to hint TorchScriptcompiler the type ofthe_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 wheretype inference can be wrong, including:

  • Empty containers like[] and{}, which TorchScript assumes to be container ofTensor

  • Optional types likeOptional[T] but assigned a valid value of typeT, TorchScript would assumeit is typeT rather thanOptional[T]

Note thatannotate() does not help in__init__ method oftorch.nn.Module subclasses because itis executed in eager mode. To annotate types oftorch.nn.Module attributes,useAttribute() instead.

Example:

importtorchfromtypingimportDict@torch.jit.scriptdeffn():# 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 forthe_value

  • the_value – Value or expression to hint type for.

Returns

the_value is passed back as return value.