Repository: pytorch/tensordict (language: python)

Issue:
[BugFix] auto-batch-size in dispatch
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #1109
* #1108
* #1107
* #1106




Discussion / hints:

Fix patch:
diff --git a/tensordict/nn/common.py b/tensordict/nn/common.py
index c38bd31de..e6d150528 100644
--- a/tensordict/nn/common.py
+++ b/tensordict/nn/common.py
@@ -301,7 +301,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
                 tensordict = make_tensordict(
                     tensordict_values,
                     batch_size=batch_size,
-                    auto_batch_size=False,
+                    auto_batch_size=self.auto_batch_size,
                 )
                 if _self is not None:
                     out = func(_self, tensordict, *args, **kwargs)


Test patch:
diff --git a/test/test_nn.py b/test/test_nn.py
index 8c78dd3b2..3134932e9 100644
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -559,6 +559,25 @@ def forward(self, tensordict):
         with pytest.raises(RuntimeError, match="Duplicated argument"):
             module(torch.zeros(1, 2), a_c=torch.ones(1, 2))
 
+    @pytest.mark.parametrize("auto_batch_size", [True, False])
+    def test_dispatch_auto_batch_size(self, auto_batch_size):
+        class MyModuleNest(nn.Module):
+            in_keys = [("a", "c"), "d"]
+            out_keys = ["b"]
+
+            @dispatch(auto_batch_size=auto_batch_size)
+            def forward(self, tensordict):
+                if auto_batch_size:
+                    assert tensordict.shape == (2, 3)
+                else:
+                    assert tensordict.shape == ()
+                tensordict["b"] = tensordict["a", "c"] + tensordict["d"]
+                return tensordict
+
+        module = MyModuleNest()
+        b = module(torch.zeros(2, 3), d=torch.ones(2, 3))
+        assert (b == 1).all()
+
     def test_dispatch_nested_extra_args(self):
         class MyModuleNest(nn.Module):
             in_keys = [("a", "c"), "d"]
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
index 27b53b686..63ebc8935 100644
--- a/test/test_tensordict.py
+++ b/test/test_tensordict.py
@@ -9865,7 +9865,7 @@ def check_weakref_count(weakref_list, expected):
 
     @pytest.mark.skipif(
         not torch.cuda.is_available(),
-     # and not torch.backends.mps.is_available(),
+        # and not torch.backends.mps.is_available(),
         reason="a device is required.",
     )
     def test_cached_data_lock_device(self):