Bug fix in metal binary kernel dispatch for large arrays (#125)

* bug fix * format
2025-09-18 18:28:12 +08:00 · 2023-12-10 16:12:31 -08:00
parent 0cfbfc9904
commit 71d1fff90a
3 changed files with 10 additions and 4 deletions
--- a/python/mlx/nn/layers/activations.py
+++ b/python/mlx/nn/layers/activations.py
@@ -22,7 +22,6 @@ def sigmoid(x):
        \text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)}
    """
    return mx.sigmoid(x)
-    


 def relu(x):
@@ -89,10 +88,12 @@ def gelu_fast_approx(x):
    """
    return x * mx.sigmoid(1.773 * x)

+
@_make_activation_module
 class Sigmoid(Module):
    pass

+
@_make_activation_module(relu)
 class ReLU(Module):
    pass