[Refactor] Update nn inline_inbuilt check

ghstack-source-id: 86c8a6dacd50387f76fd0a5b9ec9fd643b6d057f Pull Request resolved: #1029
pytorch · Oct 4, 2024 · d605e5c · d605e5c
1 parent 04faf40
commit d605e5c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 10 deletions.
diff --git a/tensordict/utils.py b/tensordict/utils.py
@@ -2507,9 +2507,9 @@ def _lock_warn():
 
 
 def _check_inbuild():
-    if not strtobool(os.environ.get("TORCHDYNAMO_INLINE_INBUILT_NN_MODULES", "0")):
+    if not torch._dynamo.config.inline_inbuilt_nn_modules:
         raise RuntimeError(
-            "to_module requires TORCHDYNAMO_INLINE_INBUILT_NN_MODULES to be set."
+            "to_module requires torch._dynamo.config.inline_inbuilt_nn_modules to be set to True."
         )
 
 

diff --git a/test/test_compile.py b/test/test_compile.py
@@ -6,7 +6,6 @@
 import contextlib
 import importlib.util
 import inspect
-import os
 from pathlib import Path
 from typing import Any, Callable
 
@@ -662,10 +661,10 @@ def test_dispatch_tensor(self, mode):
 @pytest.mark.parametrize("mode", [None, "reduce-overhead"])
 class TestFunctional:
     def test_functional_error(self, mode):
-        TORCHDYNAMO_INLINE_INBUILT_NN_MODULES = os.environ.get(
-            "TORCHDYNAMO_INLINE_INBUILT_NN_MODULES"
+        TORCHDYNAMO_INLINE_INBUILT_NN_MODULES = (
+            torch._dynamo.config.inline_inbuilt_nn_modules
         )
-        os.environ["TORCHDYNAMO_INLINE_INBUILT_NN_MODULES"] = "1"
+        torch._dynamo.config.inline_inbuilt_nn_modules = True
         module = torch.nn.Sequential(
             torch.nn.Linear(3, 4),
             torch.nn.ReLU(),
@@ -675,7 +674,7 @@ def test_functional_error(self, mode):
         td_zero = TensorDictParams(td.data.clone())
         td_zero.zero_()
 
-        os.environ["TORCHDYNAMO_INLINE_INBUILT_NN_MODULES"] = "0"
+        torch._dynamo.config.inline_inbuilt_nn_modules = False
         try:
 
             def call(x, td):
@@ -685,12 +684,12 @@ def call(x, td):
             call_compile = torch.compile(call, fullgraph=True, mode=mode)
             x = torch.randn(2, 3)
             with pytest.raises(
-                RuntimeError, match="TORCHDYNAMO_INLINE_INBUILT_NN_MODULES"
+                RuntimeError, match="torch._dynamo.config.inline_inbuilt_nn_modules"
             ):
                 call_compile(x, td_zero)
         finally:
-            if TORCHDYNAMO_INLINE_INBUILT_NN_MODULES is not None:
-                os.environ["TORCHDYNAMO_INLINE_INBUILT_NN_MODULES"] = (
+            if torch._dynamo.config.inline_inbuilt_nn_modules is not None:
+                torch._dynamo.config.inline_inbuilt_nn_modules = (
                     TORCHDYNAMO_INLINE_INBUILT_NN_MODULES
                 )