added missing tf32 support

ROCm · Nov 4, 2024 · 5fd2117 · 5fd2117
1 parent ad55450
commit 5fd2117
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tensorflow/compiler/xla/stream_executor/gpu/gpu_blas_lt.cc b/tensorflow/compiler/xla/stream_executor/gpu/gpu_blas_lt.cc
@@ -96,7 +96,8 @@ xla::StatusOr<ComputationType> GetBlasComputationType(
     case DataType::kFloat:  // fall-through
       if (lhs_dtype == DataType::kHalf) return f16_comp;
       if (lhs_dtype == DataType::kBF16) return bf16_comp;
-      return ComputationType::kF32;
+      return TF32_Enabled() ? ComputationType::kTF32AsF32 
+                            : ComputationType::kF32;
     case DataType::kComplexFloat:
       return ComputationType::kF32;
     case DataType::kDouble: // fall-through