Update on "[ET-VK][ez] Allow high dimensional tensors (for buffer storage)"

Differential Revision: [D80800083](https://our.internmc.facebook.com/intern/diff/D80800083)

[ghstack-poisoned]

Author: ssjia

backends/vulkan/runtime/api/containers/Tensor.h

@@ -705,7 +705,8 @@ std::vector<T> flip_and_unsqueeze(
     const size_t numel,
     const int32_t fixed_ndim = -1) {
   const size_t ndim = tensor_metadata.size();
-  size_t ndim_up4 = utils::align_up_4(tensor_metadata.size());
+  size_t ndim_up4 =
+      std::max(utils::align_up_4(tensor_metadata.size()), size_t(4));
 
   if (fixed_ndim > 0) {
     VK_CHECK_COND(fixed_ndim >= ndim);

backends/vulkan/runtime/graph/ops/impl/Staging.cpp

@@ -18,8 +18,6 @@
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
 
-#include <iostream>
-
 namespace vkcompute {
 
 void add_staging_to_tensor_node(