pytorch
diff --git a/‎test/dynamo/test_profiler.py‎
Lines changed: 41 additions & 0 deletions b/‎test/dynamo/test_profiler.py‎
Lines changed: 41 additions & 0 deletions
diff --git a/‎torch/_dynamo/trace_rules.py‎
Lines changed: 1 addition & 1 deletion b/‎torch/_dynamo/trace_rules.py‎
Lines changed: 1 addition & 1 deletion
diff --git a/‎torch/_dynamo/variables/torch.py‎
Lines changed: 1 addition & 0 deletions b/‎torch/_dynamo/variables/torch.py‎
Lines changed: 1 addition & 0 deletions
@@ -192,6 +192,47 @@ def fn(x, y):
             ],
         )
 
+    def test_profiler_enabled(self):
+        def fn(x):
+            x = torch.sin(x)
+            if torch.autograd._profiler_enabled():
+                return torch.cos(x)
+            else:
+                return torch.sigmoid(x)
+
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        x = torch.randn(4)
+
+        ref = fn(x)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+        with torch.autograd.profiler.profile():
+            ref = fn(x)
+            res = opt_fn(x)
+            self.assertEqual(ref, res)
+
+    def test_profiler_record_function_ignore(self):
+        def fn(x):
+            x = torch.sin(x)
+            if torch.autograd._profiler_enabled():
+                with torch.autograd.profiler.record_function("dummy"):
+                    return torch.cos(x)
+            else:
+                return torch.sigmoid(x)
+
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        x = torch.randn(4)
+
+        ref = fn(x)
+        res = opt_fn(x)
+        self.assertEqual(ref, res)
+
+        with torch.autograd.profiler.profile():
+            ref = fn(x)
+            res = opt_fn(x)
+            self.assertEqual(ref, res)
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
 
@@ -176,7 +176,6 @@
     "torch.compiler.is_compiling": TorchInGraphFunctionVariable,
     "torch.compiler.is_dynamo_compiling": TorchInGraphFunctionVariable,
     "torch.compiler.is_exporting": TorchInGraphFunctionVariable,
-    "torch.autograd._profiler_enabled": SkipFunctionVariable,
     "torch._C._to_dlpack": SkipFunctionVariable,
     "torch.to_dlpack": SkipFunctionVariable,
     # We graph break on RNG state setters or getters like
@@ -2434,6 +2433,7 @@
         "torch.atleast_3d",
         "torch.autograd._calculate_shape",
         "torch.autograd._is_checkpoint_valid",
+        "torch.autograd._profiler_enabled",
         "torch.autograd._make_grads",
         "torch.autograd._register_py_tensor_class_for_device",
         "torch.autograd._tensor_or_tensors_to_tuple",
 
@@ -142,6 +142,7 @@
     torch.cuda.is_initialized,
     torch.xpu.current_device,
     torch.xpu.is_initialized,
+    torch.autograd._profiler_enabled,
 ]
 
 constant_fold_functions = [
Original file line number	Diff line number	Diff line change
`@@ -142,6 +142,7 @@`
`142`	`142`	`torch.cuda.is_initialized,`
`143`	`143`	`torch.xpu.current_device,`
`144`	`144`	`torch.xpu.is_initialized,`
	`145`	`+ torch.autograd._profiler_enabled,`
`145`	`146`	`]`
`146`	`147`
`147`	`148`	`constant_fold_functions = [`