Hi Thomas:
Thanks for your help. I changed the for loop from this one:
for j in range(10):
  b[:, j,:] += torch.squeeze(torch.matmul(u_ji[:, j, :, :], v[j, :].reshape(16, 1)))

To this:
vv = torch.unsqueeze(v, 2).expand(32, 10, 16, 1)
b = b + torch.squeeze(torch.matmul(u_ji, vv))

and it works! Thank you so much for your help!