register nn.Module test for upsample_nearest2d (#4055)

Summary:
Pull Request resolved: #4055

Register nn.Module test for ate.upsample_nearest2d.vec

Reviewed By: jorgep31415

Differential Revision: D58907319

Author: Yujie Hui

backends/vulkan/partitioner/supported_ops.py

@@ -125,6 +125,7 @@ def __contains__(self, op):
     exir_ops.edge.aten.ones_like.default,
     exir_ops.edge.aten.zeros.default,
     exir_ops.edge.aten.zeros_like.default,
+    exir_ops.edge.aten.upsample_nearest2d.vec,
 ]
 
 

backends/vulkan/test/test_vulkan_delegate.py

@@ -1057,6 +1057,23 @@ def forward(self, x):
             memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
         )
 
+    def test_vulkan_backend_upsample_nearest2d(self):
+        class UpsampleNearest2d(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.upsample = torch.nn.Upsample(scale_factor=2, mode="nearest")
+
+            def forward(self, x):
+                return self.upsample(x)
+
+        sample_inputs = (torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2),)
+
+        self.lower_module_and_test_output(
+            UpsampleNearest2d(),
+            sample_inputs,
+            memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
+        )
+
     def test_vulkan_backend_reshape(self):
         class ReshapeModule(torch.nn.Module):
             def __init__(self):