[Inductor] No longer throw error in bmm out_dtype lowering due to template heuristics (pytorch#166457)

Fixes pytorch#165892

Pull Request resolved: pytorch#166457
Approved by: https://github.com/coconutruben

Author: PaulZhang12

torch/_inductor/kernel/bmm.py

@@ -239,9 +239,10 @@ def _to_dtype(x):
         templates_to_use.append(aten_handler)
         kwarg_overrides[aten_handler.uid] = aten_extra_kwargs
 
-    if use_triton_template(layout, check_max_autotune=False):
+    if use_triton_template(layout, check_max_autotune=False) and (
+        out_dtype is None or out_dtype == mat1.get_dtype()
+    ):
         # TODO: add out_dtype support for Triton Template
-        assert out_dtype is None, "out_dtype is not supported for Triton"
         templates_to_use.append(bmm_template)
 
     # Single unified call for all templates