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

torch.jit.isinstance(obj, target_type)[source][source]¶

Provide container type refinement in TorchScript.

It can refine parameterized containers of the List, Dict, Tuple, and Optional types. E.g. List[str],Dict[str, List[torch.Tensor]], Optional[Tuple[int,str,int]]. It can also refine basic types such as bools and ints that are available in TorchScript.

Parameters

• obj – object to refine the type of
• target_type – type to try to refine obj to

Returns

True if obj was successfully refined to the type of target_type,

False otherwise with no new type refinement

Return type

bool

Example (using torch.jit.isinstance for type refinement): .. testcode:

import torch from typing import Any, Dict, List

class MyModule(torch.nn.Module): def init(self) -> None: super().init()

def forward(self, input: Any): # note the Any type
    if torch.jit.isinstance(input, List[torch.Tensor]):
        for t in input:
            y = t.clamp(0, 0.5)
    elif torch.jit.isinstance(input, Dict[str, str]):
        for val in input.values():
            print(val)

m = torch.jit.script(MyModule()) x = [torch.rand(3,3), torch.rand(4,3)] m(x) y = {"key1":"val1","key2":"val2"} m(y)