Bug fix partitioner (#5239)

Summary:
Pull Request resolved: #5239

Forgot to copy over this part of the constraint when I refactored the partitioner

Reviewed By: balakv504, kirklandsign

Differential Revision: D62471496

fbshipit-source-id: d2fb76ffd2dea5671b38f89e050beac3a35aff4e

Author: mcr229

backends/xnnpack/partition/config/node_configs.py

@@ -85,6 +85,11 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         supported_dtypes = {torch.float32, torch.float16, torch.int8, torch.qint8}
         node_val = node.meta.get("val")
         output_0 = node_val[0]
+
+        input_node = node.all_input_nodes[0]
+        if len(input_node.meta.get("val").shape) != 4:
+            why(node, f"Unsupported input rank {input_node.meta.get('val').shape}")
+            return False
         # Don't check indicies dtype
         if output_0.dtype not in supported_dtypes:
             why(node, f"Unsupported output dtype {output_0.dtype}")