你好，我有两个张量：
a = torch.randn(125, 128) # Shape: (125, 128)
b = torch.randn(128, 8, 64) # Shape: (128, 8, 64)
我希望结果的形状为（125，8，64）
我的第一个观察是：a 的最后一个维度与 b 的第一个维度匹配，然后我这样做：
result = torch.matmul(a,b)
它给了我错误：
batch2 张量的前两个维度的预期大小为：[128, 128]，但得到的结果是：[128, 8]。
我怎样才能做到这一点。
编辑：我也不想重塑为 2D，然后再将结果重塑为 3D。
您可以使用 einsum