code mismatch

[icefree]

book/001-story-of-a-tensor/simple-network-mlx.py

def __call__(self, x):
    for i, l in enumerate(self.layers):
        # Apply ReLU to all but the last layer
        x = mx.maximum(x, 0) if i > 0 else x
        x = l(x)
    return x

the comment says "Apply ReLU to all but the last layer", but code suggests that ReLU only applies on layers with index > 0,
which is the last layer in this example

[Wankyu Choi]

class NeuralNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [
            nn.Linear(10, 20),
            nn.Linear(20, 2)
        ]

    def __call__(self, x):
        for i, l in enumerate(self.layers):
            # Apply ReLU to the first layer only
            x = mx.maximum(x, 0) if i == 0 else x
            x = l(x)
        return x

The comment should read: "Apply ReLU to the first layer only."

Fixed.

Note that "mx.maximum(x, 0) if i > 0 else x" essentially performs the ReLU operation. You can call mlx.nn.ReLU(x), but its implementation is so simple that many coders just write it out directly.

Follow the x to visualize the layers it passes through.