torch.bmm
torch.bmm(input, mat2, *, out=None) -> Tensor
Performs a batch matrix-matrix product of matrices stored in
   "input" and "mat2".
"input" and "mat2" must be 3-D tensors each containing the same
   number of matrices.
If "input" is a (b \times n \times m) tensor, "mat2" is a (b \times
   m \times p) tensor, "out" will be a (b \times n \times p) tensor.
  \text{out}_i = \text{input}_i \mathbin{@} \text{mat2}_i

This operator supports TensorFloat32.
On certain ROCm devices, when using float16 inputs this module will
   use different precision for backward.
Note:
 This function does not broadcast. For broadcasting matrix
 products, see "torch.matmul()".

Parameters:
      * input (Tensor) -- the first batch of matrices to be
        multiplied
  * **mat2** (*Tensor*) -- the second batch of matrices to be
    multiplied

Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example: