Having some trouble to vectorize these loops ([N, T, K] and [K, C] to [N, T, C])

Hello !

I’m currently having some trouble with a quite simple task, so here I am.

I have a tensor A of size [N, T, K], and a tensor B of size [K, C].
I want to compute a tensor Y of size [N, T, C] according to the following algorithm:

``````for i in range(N):
for j in range(T):
for l in range(C):
Y[i, j, l] = 0
for k in range(K):
Y[i, j, l] += A[i, j, k] * B[k, l]
``````

I’ve tried many strange thing (always ending with a “RuntimeError: the size of … must match the size of …”), but cannot find an efficient way to vectorize this.
I feel like it’s a perfect job for einsum, but didn’t find the correct notation.

`torch.matmul(A, B)` will do the job
By the way, for future reader : einsum notation is `torch.einsum('ijk,kl->ijl', (A, B))`