From 44c2ac207cd332ebe32030db87ab7a639b1f4789 Mon Sep 17 00:00:00 2001 From: vfdev Date: Fri, 10 Nov 2023 16:13:30 +0100 Subject: [PATCH 1/2] Added support for tv tensors in torch compile for func ops --- torchvision/tv_tensors/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/torchvision/tv_tensors/__init__.py b/torchvision/tv_tensors/__init__.py index d55e10e8620..2e58d9d4c6a 100644 --- a/torchvision/tv_tensors/__init__.py +++ b/torchvision/tv_tensors/__init__.py @@ -8,6 +8,10 @@ from ._video import Video +# TODO: Fix this. We skip this method as it leads to +# RecursionError: maximum recursion depth exceeded while calling a Python object +# Keeping it here, leads to graph breaks between multiple functional ops instead of having a single graph +@torch.compiler.disable def wrap(wrappee, *, like, **kwargs): """[BETA] Convert a :class:`torch.Tensor` (``wrappee``) into the same :class:`~torchvision.tv_tensors.TVTensor` subclass as ``like``. From 97bf1778105c9422dc06fc52c962d29f74894b12 Mon Sep 17 00:00:00 2001 From: vfdev Date: Fri, 10 Nov 2023 16:25:02 +0100 Subject: [PATCH 2/2] Update torchvision/tv_tensors/__init__.py Co-authored-by: Nicolas Hug --- torchvision/tv_tensors/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/tv_tensors/__init__.py b/torchvision/tv_tensors/__init__.py index 2e58d9d4c6a..441502df2f0 100644 --- a/torchvision/tv_tensors/__init__.py +++ b/torchvision/tv_tensors/__init__.py @@ -10,7 +10,7 @@ # TODO: Fix this. We skip this method as it leads to # RecursionError: maximum recursion depth exceeded while calling a Python object -# Keeping it here, leads to graph breaks between multiple functional ops instead of having a single graph +# Until `disable` is removed, there will be graph breaks after all calls to functional transforms @torch.compiler.disable def wrap(wrappee, *, like, **kwargs): """[BETA] Convert a :class:`torch.Tensor` (``wrappee``) into the same :class:`~torchvision.tv_tensors.TVTensor` subclass as ``like``.