amend

pytorch · Jul 9, 2024 · b2fa826 · b2fa826
1 parent edc315b
commit b2fa826
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 6 deletions.
diff --git a/tensordict/_lazy.py b/tensordict/_lazy.py
@@ -92,13 +92,13 @@
 _has_functorch = False
 try:
     try:
-        from torch._C._functorch import (
+        from torch._C._functorch import (  # @manual=fbcode//caffe2:_C
             _add_batch_dim,
             _remove_batch_dim,
             is_batchedtensor,
         )
     except ImportError:
-        from functorch._C import is_batchedtensor
+        from functorch._C import is_batchedtensor  # @manual=fbcode//caffe2/functorch:_C
 
     _has_functorch = True
 except ImportError:

diff --git a/tensordict/_td.py b/tensordict/_td.py
@@ -101,13 +101,13 @@
 _has_functorch = False
 try:
     try:
-        from torch._C._functorch import (
+        from torch._C._functorch import (  # @manual=fbcode//caffe2:_C
             _add_batch_dim,
             _remove_batch_dim,
             is_batchedtensor,
         )
     except ImportError:
-        from functorch._C import is_batchedtensor
+        from functorch._C import is_batchedtensor  # @manual=fbcode//functorch:_C
 
     _has_functorch = True
 except ImportError:

diff --git a/tensordict/utils.py b/tensordict/utils.py
@@ -69,9 +69,15 @@
 
 try:
     try:
-        from functorch._C import get_unwrapped, is_batchedtensor
+        from torch._C._functorch import (  # @manual=fbcode//caffe2:_C
+            get_unwrapped,
+            is_batchedtensor,
+        )
     except ImportError:
-        from torch._C._functorch import get_unwrapped, is_batchedtensor
+        from functorch._C import (
+            get_unwrapped,
+            is_batchedtensor,
+        )  # @manual=fbcode//functorch:_C
 except ImportError:
     pass