Update on "fix issue of baddbmm when out has nan value for beta=0"

Fix #96037.

cc jgong5 mingfeima sanchitintel ashokei jingxu10

[ghstack-poisoned]

aten/src/ATen/native/LinearAlgebra.cpp

@@ -1557,6 +1557,7 @@ inline void baddbmm_cpu_kernel(const Tensor& result, const Tensor& self, const T
                 r += s2[k] * m1[k][j];
               }
             } else {
+              // For beta == 0, the r's value will be ignored, especially for nan value.
               r = beta == scalar_t(0) ? scalar_t(0) : beta * r;
               for (const auto k : c10::irange(ks)) {
                 r += alpha * s2[k] * m1[k][j];