Multiply a tensor of 2D matrices with a (batched) tensor of vectors

Given these tensors:

batch_size = 5
F = torch.randn(k+1, 3, 3)
B = torch.randn(batch_size, k, 1)

I’m trying to perform the following operation but without using the for loop (with is pretty slow when A and B are huge tensors):

for i in range(1,F.shape[0]-1):
	S += F[i,:,:]*B[:,i:(i+1),0:1]

Is it possible to do this without using the for loop? I am aware of the broadcasting capabilities of Pytorch, but I am not sure how I could apply it to this example.


I just solved it using einsum (see tutorial here), it is way faster than the loop

S=torch.einsum('ajk,ial->ijk', [F[0:-1,:,:], B])