Diag a demension of a matrix

Hi all,
I have a matrix A(A.shape=torch.Size([50, 400]))and a batch of vector B(B.shape=torch.Size([8,400])),where 8 is the batch size.
The first thing I want to do is make the vector part of B (a.k.a [400]) a diagonal matrix, then B should become a matrix in shape [8, 400, 400]
The I want to perform patch multiplication, that is calculate [50, 400] multiply [8,400,400] and get the result in size [8,50,400].

I want to know how to transform some dimension in a matrix like B into diagonal matrix(like from[8,400] to [8,400,400])
I want to do these things in a neural work so if the network can still be backpropagated after I converted second dimension of B into a diagonal matrix? Or should I implement the same operation on A and B with another way that don’t need diag operation?

Question closed, it is not neccessary to diag a vector to a diagonal matrix, you can achieve element-wise multiply simply by using mul()
and torch will broadcast for you.
a=[8, 1, 400]
b=[50, 400]
just use a.mul(b)
ans is [8, 50, 400]