Add an option to disable aten::cat in TE (re-revert) (#50101)

Summary:
This reverts commit ace78dd.

Fixes #{issue number}

Pull Request resolved: #50101

Reviewed By: eellison

Differential Revision: D25784785

Pulled By: Krovatkin

fbshipit-source-id: cbb3d377e03303f6c8c71f4c59c6d90ab40d55f7

torch/csrc/jit/passes/tensorexpr_fuser.cpp

@@ -15,6 +15,11 @@
 #include <torch/csrc/jit/tensorexpr/kernel.h>
 #include <torch/csrc/utils/memory.h>
 
+// NOLINTNEXTLINE
+C10_DEFINE_bool(
+    torch_jit_disable_cat,
+    false,
+    "disable aten::cat in TE fusion groups");
 namespace torch {
 namespace jit {
 
@@ -202,6 +207,10 @@ bool isSupported(Node* node) {
       }
     }
 
+    if (FLAGS_torch_jit_disable_cat && node->kind() == aten::cat) {
+      return false;
+    }
+
     return true;
   }