et_converter: Identify non-comm nodes as comp nodes

mlcommons · Feb 6, 2024 · b56f0b8 · b56f0b8
1 parent 0da8aa6
commit b56f0b8
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/et_converter/pytorch2chakra_converter.py b/et_converter/pytorch2chakra_converter.py
@@ -524,13 +524,9 @@ def get_chakra_node_type_from_pytorch_node(self, pytorch_node: PyTorchNode) -> i
             "ncclKernel" in pytorch_node.name or "ncclDevKernel" in pytorch_node.name
         ):
             return COMM_COLL_NODE
-        elif pytorch_node.is_gpu_op():
-            return COMP_NODE
         elif ("c10d::" in pytorch_node.name) or ("nccl:" in pytorch_node.name):
             return COMM_COLL_NODE
-        elif (pytorch_node.op_schema != "") or pytorch_node.outputs:
-            return COMP_NODE
-        return INVALID_NODE
+        return COMP_NODE
 
     def get_collective_comm_type(self, name: str) -> int:
         """