[fbsync] Added support for tv tensors in torch compile for func ops (#…

…8110) Summary: Co-authored-by: Nicolas Hug <[email protected]> Reviewed By: vmoens Differential Revision: D51391968 fbshipit-source-id: 90b6b67404f5faac9ff4effcf3d839f77fc9f7b1
pytorch · Nov 16, 2023 · 856e7cf · 856e7cf
1 parent 0f588e4
commit 856e7cf
Showing 1 changed file with 4 additions and 0 deletions.
diff --git 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
+# 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``.