PyTorch - Multiply 1D tensor across dimension 0 of 5D tensor

Question

So I have this tensor A that is 1-dimensional of length 6 then I have another tensor B that's 5-dimensional and dimension 0 is of shape (6, batch size, 3, 100, 100) . I'm trying to multiply all values in B[n,:] (4-dimensional tensor of shape (batch size, 3, 100, 100)) by the scalar A[n] where n is between 0 and 5. How can I do this? The operation should also be differentiable because A comes out of a neural network and the 5-dimensional tensor after this operation goes into another neural network which I'd like to backpropagate on fully.

I've tried this

inputs = A.view(6, 1, 1, 1, 1) * B

But it doesn't seem like the gradient flows through the neural network that outputs A.