From 608f2478fc945b90d87f91442fb5474d01e6ffde Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Tue, 3 Dec 2024 11:45:38 +0000
Subject: [PATCH] Update (base update)

[ghstack-poisoned]
---
 tensordict/tensorclass.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensordict/tensorclass.py b/tensordict/tensorclass.py
index 54d4f3b9c..e5a433af2 100644
--- a/tensordict/tensorclass.py
+++ b/tensordict/tensorclass.py
@@ -123,6 +123,8 @@ def __subclasscheck__(self, subclass):
 ]
 # Methods to be executed from tensordict, any ref to self means 'self._tensordict', no wrap of result
 _FALLBACK_METHOD_FROM_TD_NOWRAP = [
+    "_check_batch_size",
+    "_check_device",
     "_check_dim_name",
     "_check_unlock",
     "_default_get",
@@ -181,6 +183,7 @@ def __subclasscheck__(self, subclass):
     "_select",  # TODO: must be specialized
     "_set_at_tuple",
     "_set_tuple",
+    "_to_module",
     "abs",
     "abs_",
     "acos",
@@ -327,6 +330,7 @@ def __subclasscheck__(self, subclass):
     "tanh",
     "tanh_",
     "to",
+    "to_module",
     "transpose",
     "trunc",
     "trunc_",