[MPS] Fix embedding cache key

ghstack-source-id: 3ba9c3d6cc9cfa7977cbe43f412dad099aef5bad Pull Request resolved: #101857
pytorch · May 19, 2023 · 6b85dd2 · 6b85dd2
1 parent f994d0b
commit 6b85dd2
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/aten/src/ATen/native/mps/operations/Indexing.mm b/aten/src/ATen/native/mps/operations/Indexing.mm
@@ -819,9 +819,8 @@ Tensor embedding_dense_backward_mps(const Tensor& grad_,
   auto stream = at::mps::getCurrentMPSStream();
 
   @autoreleasepool {
-    string key = "edb_mps:" + getMPSTypeString(grad_) + ":indices" + std::to_string(num_indices_dims) + ":num_weights" +
-        std::to_string(num_weights) + ":padding_idx" + std::to_string(padding_idx) + ":scaled" +
-        std::to_string(scale_grad_by_freq);
+    string key = "edb_mps:" + getTensorsStringKey({grad_, indices}) + ":num_weights" + std::to_string(num_weights) +
+        ":padding_idx" + std::to_string(padding_idx) + ":scaled" + std::to_string(scale_grad_by_freq);
     auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
       MPSGraphTensor* incomingGradTensor = mpsGraphUnrankedPlaceHolder(mpsGraph, getMPSDataType(grad_));
 

diff --git a/test/test_mps.py b/test/test_mps.py
@@ -6719,6 +6719,7 @@ def helper(n, d, m, idx):
             self.assertEqual(a_CPU.grad, a_MPS.grad)
 
         helper(3, 5, 7, [0, 1, 2])
+        helper(3, 6, 7, [0, 1, 2])
         helper(3, 5, 7, 2)  # test scalar index
 
     # Test pytorch gather