set_return_type

class torchvision.tv_tensors.set_return_type(return_type: str)

Set the return type of torch operations on TVTensor.

This only affects the behaviour of torch operations. It has no effect on torchvision transforms or functionals, which will always return as output the same type that was passed as input.

Warning

We recommend using ToPureTensor at the end of your transform pipelines if you use set_return_type("TVTensor"). This will avoid the __torch_function__ overhead in the models forward().

Can be used as a global flag for the entire program:

img = tv_tensors.Image(torch.rand(3, 5, 5))
img + 2  # This is a pure Tensor (default behaviour)

set_return_type("TVTensor")
img + 2  # This is an Image

or as a context manager to restrict the scope:

img = tv_tensors.Image(torch.rand(3, 5, 5))
img + 2  # This is a pure Tensor
with set_return_type("TVTensor"):
    img + 2  # This is an Image
img + 2  # This is a pure Tensor

Parameters:

return_type (str) – Can be “TVTensor” or “Tensor” (case-insensitive). Default is “Tensor” (i.e. pure torch.Tensor).

Examples using set_return_type:

TVTensors FAQ

TVTensors FAQ