Simplify check_node_has_valid_dtype (#1801)

Summary:
Pull Request resolved: #1801

Linter seems to be unhappy.

https://github.com/pytorch/executorch/actions/runs/7736952506/job/21095126521

Reviewed By: mcr229

Differential Revision: D53313246

fbshipit-source-id: 1aef8d38a6816db8292ec25e0c1bce5db0a545ef

Author: mergennachin

backends/xnnpack/partition/xnnpack_partitioner.py

@@ -100,22 +100,7 @@ def __init__(
         self.ep = ep
         assert len(self.constraints)
 
-    def check_node_has_valid_dtype(self, node):
-        if node.target in {exir_ops.edge.aten.max_pool2d_with_indices.default}:
-            return True
-
-        valid_dtypes = {
-            torch.float32,
-            torch.int8,
-            torch.qint8,
-        }
-        if (
-            node.op != "placeholder"
-            and node.op != "call_function"
-            and node.op != "get_attr"
-        ):
-            return False
-
+    def _check_inputs_are_valid_dtypes(self, node, valid_dtypes):
         # Check inputs are valid dtypes
         for arg in node.args:
             if not isinstance(arg, torch.fx.Node):
@@ -133,21 +118,46 @@ def check_node_has_valid_dtype(self, node):
             if arg_val.dtype not in valid_dtypes:
                 return False
 
+        return True
+
+    def _check_outputs_are_valid_dtypes(self, node, valid_dtypes):
         # Check outputs are valid dtype
         node_val = node.meta.get("val", None)
-        if node_val is not None:
-            if not isinstance(node_val, tuple):
-                node_val = (node_val,)
+        if node_val is None:
+            return True
 
-            for val in node_val:
-                if not isinstance(val, torch.Tensor):
-                    return False
+        if not isinstance(node_val, tuple):
+            node_val = (node_val,)
 
-                if val.dtype not in valid_dtypes:
-                    return False
+        for val in node_val:
+            if not isinstance(val, torch.Tensor):
+                return False
+
+            if val.dtype not in valid_dtypes:
+                return False
 
         return True
 
+    def check_node_has_valid_dtype(self, node):
+        if node.target in {exir_ops.edge.aten.max_pool2d_with_indices.default}:
+            return True
+
+        valid_dtypes = {
+            torch.float32,
+            torch.int8,
+            torch.qint8,
+        }
+        if (
+            node.op != "placeholder"
+            and node.op != "call_function"
+            and node.op != "get_attr"
+        ):
+            return False
+
+        return self._check_inputs_are_valid_dtypes(
+            node, valid_dtypes
+        ) and self._check_outputs_are_valid_dtypes(node, valid_dtypes)
+
     def check_common_constraints(self, node) -> bool:
         has_valid_dtypes = self.check_node_has_valid_dtype(node)