我正在尝试编写一个基于计算嵌入的Pytorch程序，在训练过程中使用嵌入的各种子集。我想使用Tensor索引对嵌入Tensor进行子集化，但是当我这样做时，我会得到典型的向后二次误差。但是，使用切片而不是索引来执行完全相同的操作可以运行得很好。在我的实际代码中，我已经确认了切片版本可以按预期工作并正确训练等，但新功能不能只用切片来完成。
我在Windows上运行python 3.10.9，pytorch 2.0.0和cuda 11.8。
下面是一个最小的例子，重现了我遇到的错误：

import torch
device = 'cpu'    # same error occurs when using device='cuda:0'
embeddings = torch.tensor(torch.randn([1024, 128], device=device, dtype=torch.float32), requires_grad=True)
target = torch.rand([1024, 128], device=device)
optimizer = torch.optim.Adadelta([embeddings], lr=1.0)

# alternate approach using nn.Embedding doesn't work either
embeddings_2 = torch.nn.Embedding(1024, 128, device=device)

# OPTION 1: This works
# cur_embeddings = embeddings[:512, :]

# OPTION 2: This is what I want to do, but it doesn't work
cur_embeddings = embeddings[torch.arange(512, device=device), :]

# OPTION 3: The following option does not work either
# cur_embeddings = embeddings_2(torch.arange(512, device=device))

cur_target = target[:512, :]
for idx in range(2):
    optimizer.zero_grad()
    loss = torch.nn.MSELoss()(cur_embeddings, cur_target)
    loss.backward()
    optimizer.step()

在所呈现的3个选项中，选项1工作，但选项2产生错误，即使其功能相同。如何使选项2的工作方式与选项1完全相同（同时保持对任意行进行子集划分的自由）？
我还在这里展示了使用nn.Embedding（我并不真的想使用它，但如果它能工作的话，我可以使用它）仍然会产生错误。