[torch.compile] Compile CustomOp.forward_native for ReLU^2 to avoid raw torch ops inside opaque custom ops (#50244)

Signed-off-by: Roi Koren <roik@nvidia.com>
2026-07-29 18:27:31 +00:00
parent 5c7a7f9462
commit 93477454f9
1 changed files with 2 additions and 2 deletions
@@ -658,8 +658,8 @@ class ReLUSquaredActivation(CustomOp):

    # --8<-- [end:relu2]

-    def __init__(self):
-        super().__init__()
+    def __init__(self, *, compile_native: bool = True):
+        super().__init__(compile_native=compile_native)
        if current_platform.is_cuda_alike():
            self.op = torch.ops._C.relu_squared