Update on "[Bugfix] Allow non Module in TensorDictModule when method is passed"

relax this check as [LLM](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/llm.py#L53) in vLLM does not subclass nn.Module




[ghstack-poisoned]

Author: mikaylagawarecki

tensordict/nn/common.py

@@ -1072,14 +1072,12 @@ def __init__(
         if method is None:
             if type(module) is type or not callable(module):
                 raise ValueError(
-                    f"Module {module} if type {type(module)} is not callable. "
+                    f"Module {module} of type {type(module)} is not callable. "
                     f"Typical accepted types are nn.Module or TensorDictModule."
                 )
         else:
             if not (hasattr(module, method) and callable(getattr(module, method))):
-                raise ValueError(
-                    f"Module {module} does not have a callable method {method}. "
-                )
+                raise ValueError(f"Module {module} does not have a method {method}. ")
         self.out_keys = out_keys
         self.in_keys = in_keys
 

test/test_nn.py

@@ -295,18 +295,21 @@ def my_func(self, tensor: torch.Tensor, *, an_integer: int):
         td = s(TensorDict(a=0))
 
         assert td["b"] == 4
-    
+
     def test_non_module_method(self):
-        class MyNet:
+        class NonModuleNet:
+            def __init__(self):
+                self.foo = 2
+
             def my_func(self, tensor: torch.Tensor, *, an_integer: int):
                 return tensor + an_integer
-        
+
         s = TensorDictSequential(
             {
                 "a": lambda td: td + 1,
                 "b": lambda td: td * 2,
                 "c": TensorDictModule(
-                    MyNet(),
+                    NonModuleNet(),
                     in_keys=["a"],
                     out_keys=["b"],
                     method="my_func",
@@ -319,6 +322,15 @@ def my_func(self, tensor: torch.Tensor, *, an_integer: int):
 
         assert td["b"] == 4
 
+        with pytest.raises(ValueError, match="does not have a method foo."):
+            TensorDictModule(
+                NonModuleNet(),
+                in_keys=["a"],
+                out_keys=["b"],
+                method="foo",
+                method_kwargs={"an_integer": 2},
+            ),
+
     def test_mutable_sequence(self):
         in_keys = self.MyMutableSequence(["a", "b", "c"])
         out_keys = self.MyMutableSequence(["d", "e", "f"])