Fix typo in comment

sonos · Dec 11, 2024 · 1c85939 · 1c85939
1 parent c22bb42
commit 1c85939
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/metal/src/kernels/nn/nn_ops.metal b/metal/src/kernels/nn/nn_ops.metal
@@ -303,7 +303,7 @@ template<typename F>
     size_t mask_base_idx = tgpig.y * mask_strides[1] 
             + tgpig.z * mask_strides[0];
 
-    // Get max value on softmax reduce_dim after apply
+    // Get max value on softmax reduce_dim after applying scale and mask
     float partial_max = -INFINITY;
     for (size_t i = tiisg; i < reduce_dim; i += tpsg) {
         auto idx = base_idx + i * strides[2];