[Feature] give a __name__ to TDModules

ghstack-source-id: 95a9f4115ea45814b174484ae113085634c3d2a2
Pull Request resolved: #1045

tensordict/nn/common.py

@@ -745,6 +745,14 @@ def _reset_parameters(self, module: nn.Module) -> bool:
                 any_reset |= True
         return any_reset
 
+    @property
+    def __name__(self):
+        # This is necessary to make compiled vmap over TDModule happy
+        return self.__class__.__name__
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}()"
+
 
 class TensorDictModule(TensorDictModuleBase):
     """A TensorDictModule, is a python wrapper around a :obj:`nn.Module` that reads and writes to a TensorDict.