From 581b699ac98015876710f9da950fe3bc70cdf26d Mon Sep 17 00:00:00 2001 From: Piotr Rybiec Date: Wed, 17 Apr 2024 17:26:22 +0200 Subject: [PATCH] avgpool, not maxpool (#1002) --- python/mlx/nn/layers/pooling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/mlx/nn/layers/pooling.py b/python/mlx/nn/layers/pooling.py index 1bcd2a2dc..3733fd777 100644 --- a/python/mlx/nn/layers/pooling.py +++ b/python/mlx/nn/layers/pooling.py @@ -317,7 +317,7 @@ class AvgPool2d(_Pool2d): >>> import mlx.core as mx >>> import mlx.nn.layers as nn >>> x = mx.random.normal(shape=(8, 32, 32, 4)) - >>> pool = nn.MaxPool2d(kernel_size=2, stride=2) + >>> pool = nn.AvgPool2d(kernel_size=2, stride=2) >>> pool(x) """