Forward fix Fix batch norm partitioning with Conv3d (pytorch#13696) (pytorch#14170)

Summary:
Forward pyre fix

https://www.internalfb.com/diff/D81069236
pytorch#13696

bypass-github-export-checks
bypass-github-pytorch-ci-checks
bypass-github-executorch-ci-checks

Reviewed By: digantdesai

Differential Revision: D82118651

Author: mergennachin

backends/xnnpack/_passes/fuse_batch_norm.py

@@ -115,8 +115,14 @@ def can_fuse(
             return False
 
         # Check the rank of the convolutution input - only Conv1d and 2d are supported.
-        if is_conv and len(input_node.args[0].meta["val"].shape) not in (3, 4):
-            return False
+        if is_conv:
+            conv_input = input_node.args[0]
+            if (
+                not isinstance(conv_input, torch.fx.Node)
+                or "val" not in conv_input.meta
+                or len(conv_input.meta["val"].shape) not in (3, 4)
+            ):
+                return False
 
         return True
 

backends/xnnpack/test/passes/test_batch_norm_fusion.py

@@ -60,7 +60,7 @@ def __init__(
             self,
             in_features: int,
             out_features: int,
-            kernel_size: Tuple[int, int],
+            kernel_size: Tuple[int, int, int],
         ):
             super().__init__()
             op = torch.nn.Conv3d