在pytorch中通过乘法扩展矩阵维数

Question

我正在使用 pytorch，我有维度 [a，b，c] 的张量 A 和维度 [a，d] 的张量 B。我想创建维度 [a、b、c、d] 的张量 C，即 A 中的元素与 B 中的元素相乘，操作如下：

for i in range(a):
    for j in range(b):
        for k in range(c):
            for l in range(d):
                C[i,j,k,l]=A[i,j,k]*B[i,l]

这按预期工作，但速度很慢。在 pytorch 中进行此类操作的最佳实践是什么？

谢谢你。

Answer 1

A.reshape(a, b, c, 1) * B.reshape(a, 1, 1, d)

会成功的。