Disable type inference flag is not being set (#455)

* Disable type inference flag is not being set * Fix softmax
kulinseth · Jun 26, 2023 · 302584f · 302584f
1 parent e6f4876
commit 302584f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/aten/src/ATen/native/mps/OperationUtils.mm b/aten/src/ATen/native/mps/OperationUtils.mm
@@ -26,7 +26,7 @@ void runMPSGraph(MPSStream* mpsStream, MPSCachedGraph *cachedGraph, NSDictionary
           shapes[graphTensor] = [[[MPSGraphShapedType alloc] initWithShape:nil dataType:graphTensorData.dataType] autorelease];
         }
         MPSGraphCompilationDescriptor *compilationDescriptor = [[MPSGraphCompilationDescriptor new] autorelease];
-
+        [compilationDescriptor disableTypeInference];
         executable = [[mpsGraph compileWithDevice:nil
                                             feeds:shapes
                                     targetTensors:[results allKeys]

diff --git a/aten/src/ATen/native/mps/operations/SoftMax.mm b/aten/src/ATen/native/mps/operations/SoftMax.mm
@@ -102,7 +102,7 @@ void get_shapes(MPSShape* input_shape_readonly,
           assert(0 && "Invalid dim\n");
       }
     }
-
+    bool disableTypeInference = input_.dim() <= 4;
     string key = "softmax_mps_out" + getTensorsStringKey(input, true, /*exclude_shape*/true) + ":" +
                  mem_format_key + ":" + std::to_string(dim_);
     CachedGraph* cachedGraph = static_cast<CachedGraph *>(cache_->LookUp(key));
@@ -162,7 +162,7 @@ void get_shapes(MPSShape* input_shape_readonly,
       outputPlaceholder.getMPSGraphTensor() : outputPlaceholder.getMPSGraphTensorData()
     };
 
-    runMPSGraph(stream, cachedGraph, feeds, results, /*disable_type_inference*/ true);
+    runMPSGraph(stream, cachedGraph, feeds, results, disableTypeInference);
   }
 
 }