[MPS] Fix `test_embedding_dense_backward` (#88847)
By copying randomly initialized weights distribution from MPS `nn.Embedding` to `cpu`
Test plan: `python test_mps.py -k test_embedding_dense_backward --repeat 150`
Fixes https://github.com/pytorch/pytorch/issues/88679
Pull Request resolved: https://github.com/pytorch/pytorch/pull/88847
Approved by: https://github.com/seemethere
diff --git a/test/test_mps.py b/test/test_mps.py
index 2ff5a9d..30546f5 100644
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -4282,8 +4282,9 @@
def test_embedding_dense_backward(self):
def helper(n, d, m, idx):
embeddingMPS = nn.Embedding(n, d, max_norm=True, device='mps')
+ emedding_weight = embeddingMPS.weight.detach().cpu()
W_MPS = torch.randn((m, d), requires_grad=True, device='mps')
- idx_MPS = torch.tensor(idx).to('mps')
+ idx_MPS = torch.tensor(idx, device='mps')
a_MPS = embeddingMPS.weight.clone() @ W_MPS.t() # weight must be cloned for this to be differentiable
a_MPS.retain_grad()
b_MPS = embeddingMPS(idx_MPS) @ W_MPS.t() # modifies weight in-place
@@ -4292,7 +4293,7 @@
loss_MPS = out_MPS.sigmoid().prod()
loss_MPS.backward()
- embeddingCPU = nn.Embedding(n, d, max_norm=True, scale_grad_by_freq=True)
+ embeddingCPU = nn.Embedding(n, d, max_norm=True, _weight=emedding_weight)
W_CPU = W_MPS.to('cpu')
idx_CPU = torch.tensor(idx)
a_CPU = embeddingCPU.weight.clone() @ W_CPU.t() # weight must be cloned for this to be differentiable
