Simplify non-skipped tensor logic

Author: abhudev

aten/src/ATen/native/mps/operations/Shape.mm

@@ -545,13 +545,16 @@ void check_shape_except_dim(const Tensor &first, const Tensor &second,
 
   // Indices of tensors to be skipped because they're empty
   std::vector<int64_t> skipped_tensor_indices;
+  // Tensors to be read
+  std::vector<const Tensor*> input_tensors;
   int tensor_idx = 0;
   for(const Tensor& t : materialized_inputs) {
     if(t.numel() == 0 || should_skip(t)) {
       skipped_tensor_indices.push_back(tensor_idx);
       tensor_idx++;
       continue;
     }
+    input_tensors.push_back(&t);
     nDims = t.dim();
     // TODO: Is this OK?
     notSkippedTensor = &t;
@@ -659,26 +662,22 @@ void check_shape_except_dim(const Tensor &first, const Tensor &second,
           MPSGraphTensor* castInputMPSGraphTensors[len_tensor_array];
 
           int graph_tensor_idx = 0;
-          int i = 0;
-          for(const Tensor& tensor : materialized_inputs) {
-            if(std::find(skipped_tensor_indices.begin(), skipped_tensor_indices.end(), i) == skipped_tensor_indices.end()) {
-              inputMPSGraphTensors[graph_tensor_idx] = mpsGraphUnrankedPlaceHolder(mpsGraph, getMPSDataType(tensor.scalar_type()) );
-              if(getMPSDataType(result_type(inputs)) == MPSDataTypeBool) {
+          for(const Tensor* tensor : input_tensors) {
+            inputMPSGraphTensors[graph_tensor_idx] = mpsGraphUnrankedPlaceHolder(mpsGraph, getMPSDataType(tensor->scalar_type()) );
+            if(getMPSDataType(result_type(inputs)) == MPSDataTypeBool) {
+              castInputMPSGraphTensors[graph_tensor_idx] = [mpsGraph castTensor:inputMPSGraphTensors[graph_tensor_idx]
+                                                                           toType:MPSDataTypeFloat32
+                                                                             name:[NSString stringWithFormat:@"castInput%@", [NSNumber numberWithInt:graph_tensor_idx]]];
+            }
+            else {
+              if(tensor->scalar_type() != result_type(inputs))
                 castInputMPSGraphTensors[graph_tensor_idx] = [mpsGraph castTensor:inputMPSGraphTensors[graph_tensor_idx]
-                                                                             toType:MPSDataTypeFloat32
-                                                                               name:[NSString stringWithFormat:@"castInput%@", [NSNumber numberWithInt:graph_tensor_idx]]];
-              }
-              else {
-                if(tensor.scalar_type() != result_type(inputs))
-                  castInputMPSGraphTensors[graph_tensor_idx] = [mpsGraph castTensor:inputMPSGraphTensors[graph_tensor_idx]
-                                                                             toType:getMPSDataType(result_type(inputs))
-                                                                               name:[NSString stringWithFormat:@"castInput%@", [NSNumber numberWithInt:graph_tensor_idx]]];
-                else
-                  castInputMPSGraphTensors[graph_tensor_idx] = inputMPSGraphTensors[graph_tensor_idx];
-              }
-              graph_tensor_idx++;
+                                                                           toType:getMPSDataType(result_type(inputs))
+                                                                             name:[NSString stringWithFormat:@"castInput%@", [NSNumber numberWithInt:graph_tensor_idx]]];
+              else
+                castInputMPSGraphTensors[graph_tensor_idx] = inputMPSGraphTensors[graph_tensor_idx];
             }
-            i++;
+            graph_tensor_idx++;
           }
 
           auto inputTensorsArray = [NSArray arrayWithObjects:castInputMPSGraphTensors