[atem] Fix type check bug in bmm_out_or_baddbmm_ (#51248)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51248
Fix bug raised in https://github.com/pytorch/pytorch/issues/50980 by adding a dtype check back to bmm_out_or_baddbmm_.
Test Plan:
```
buck test //caffe2/test:linalg
buck test //caffe2/aten:math_kernel_test
```
Reviewed By: ngimel
Differential Revision: D26113575
fbshipit-source-id: 0d6e03eae70822f8ceeffefd915aee01030304ce
