amend

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

Hide file tree

Showing 4 changed files with 8 additions and 8 deletions.
diff --git a/tensordict/_lazy.py b/tensordict/_lazy.py
@@ -92,7 +92,7 @@
 _has_functorch = False
 try:
     try:
-        from torch._C._functorch import (  # @manual=fbcode//caffe2:_C
+        from torch._C._functorch import (  # @manual=fbcode//caffe2:torch
             _add_batch_dim,
             _remove_batch_dim,
             is_batchedtensor,

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

diff --git a/tensordict/nn/functional_modules.py b/tensordict/nn/functional_modules.py
@@ -109,8 +109,8 @@ def set_tensor_dict(  # noqa: F811
 
 _RESET_OLD_TENSORDICT = True
 try:
-    import torch._functorch.vmap as vmap_src
-    from torch._functorch.vmap import (
+    import torch._functorch.vmap as vmap_src  # @manual=fbcode//caffe2:torch
+    from torch._functorch.vmap import (  # @manual=fbcode//caffe2:torch
         _add_batch_dim,
         _broadcast_to_and_flatten,
         _get_name,
@@ -124,7 +124,7 @@ def set_tensor_dict(  # noqa: F811
     _has_functorch = True
 except ImportError:
     try:
-        from functorch._src.vmap import (
+        from functorch._src.vmap import (  # @manual=fbcode//caffe2/functorch:functorch_src
             _add_batch_dim,
             _broadcast_to_and_flatten,
             _get_name,
@@ -136,7 +136,7 @@ def set_tensor_dict(  # noqa: F811
         )
 
         _has_functorch = True
-        import functorch._src.vmap as vmap_src
+        import functorch._src.vmap as vmap_src  # @manual=fbcode//caffe2/functorch:functorch_src
     except ImportError:
         _has_functorch = False
 

diff --git a/tensordict/utils.py b/tensordict/utils.py
@@ -69,7 +69,7 @@
 
 try:
     try:
-        from torch._C._functorch import (  # @manual=fbcode//caffe2:_C
+        from torch._C._functorch import (  # @manual=fbcode//caffe2:torch
             get_unwrapped,
             is_batchedtensor,
         )
@@ -2321,7 +2321,7 @@ class _add_batch_dim_pre_hook:
     def __call__(self, mod: torch.nn.Module, args, kwargs):
         for name, param in list(mod.named_parameters(recurse=False)):
             if hasattr(param, "in_dim") and hasattr(param, "vmap_level"):
-                from torch._C._functorch import _add_batch_dim
+                from torch._C._functorch import _add_batch_dim  # @manual=//caffe2:_C
 
                 param = _add_batch_dim(param, param.in_dim, param.vmap_level)
                 delattr(mod, name)