From 080f02ac7a31ee0ab039b02eed08d3964e648a15 Mon Sep 17 00:00:00 2001 From: Richeek Das Date: Thu, 10 Oct 2024 19:11:37 +0000 Subject: [PATCH] [dynamo] do not raise an unimplemented error with boolean masking setitem (#134902) Cudagraph breaks on boolean masking setitem, however the code runs fine. There is no need to raise an unimplemented error here, since it already warns that its an incompatible op. Fixes #134241 Pull Request resolved: https://github.com/pytorch/pytorch/pull/134902 Approved by: https://github.com/jansel, https://github.com/henrylhtsang --- test/dynamo/test_export.py | 5 +---- torch/_dynamo/variables/tensor.py | 9 --------- 2 files changed, 1 insertion(+), 13 deletions(-) diff --git a/test/dynamo/test_export.py b/test/dynamo/test_export.py index 03d40e377335b..4fd829508f36f 100644 --- a/test/dynamo/test_export.py +++ b/test/dynamo/test_export.py @@ -4565,10 +4565,7 @@ def forward(self, x, b, y): return pytree.tree_unflatten([x], self._out_spec)""", # NOQA: B950 ) - with self.assertRaisesRegex( - torch._dynamo.exc.Unsupported, "boolean masking setitem backwards" - ): - gm, _ = torch._dynamo.export(fn)(x, b, y) + gm, _ = torch._dynamo.export(fn)(x, b, y) def test_dynamo_list_index(self): def fn(x, in_list): diff --git a/torch/_dynamo/variables/tensor.py b/torch/_dynamo/variables/tensor.py index b96d35744e73e..f0ae6cdceadd8 100644 --- a/torch/_dynamo/variables/tensor.py +++ b/torch/_dynamo/variables/tensor.py @@ -874,15 +874,6 @@ def has_bool_key(v): else: return False - if ( - has_bool_key(key) - and isinstance(value, TensorVariable) - and value.requires_grad - and torch.is_grad_enabled() - ): - unimplemented( - "boolean masking setitem backwards, see https://github.com/pytorch/pytorch/issues/114123" - ) from ..symbolic_convert import InstructionTranslator tx = InstructionTranslator.current_tx()