[PyTorch] Add flop count for bmm and baddbmm (#66636)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/66636
Add FLOP count for bmm and baddbmm, which is `2*b*m*n*k`.
Reviewed By: ngimel
Differential Revision: D31622061
fbshipit-source-id: f3e1e1e34c45228693117b81647fb4a623c4085b
