Train a Classifier

Build and train a neural network that separates data a linear model cannot.

You are given training and validation tensors. Return a callable that takes an (n, d) float tensor and returns predicted integer class labels of shape (n,). What you return can be a module, a closure, or anything with a __call__; only its behaviour is graded.

The data is two interleaving half moons, so a straight line cannot do better than roughly 85%. You need at least one hidden layer with a nonlinearity. A reasonable shape is Linear(d, 32) -> ReLU -> Linear(32, 2) trained with Adam on CrossEntropyLoss.

Two things worth knowing:

• nn.CrossEntropyLoss takes raw logits, not probabilities. Putting a softmax before it applies the log twice and the model trains badly while still appearing to work.
• Predicting means taking argmax over the class axis of those logits, and doing it under torch.no_grad().

This is scored on the result: your classifier has to reach at least 92% accuracy on the held-out validation set, which it never sees during grading beyond what you do with it yourself.

Function signature

def train_classifier(X_train, y_train):

Example

Input

predict = train_classifier(X_train, y_train)    # two moons, 2 features
(predict(X_val) == y_val).float().mean()

Output

tensor(0.97)    # must be at least 0.92

Reasoning

The two moons interleave, so a straight line gets about 85%. A hidden layer with ReLU can bend the boundary around both moons, and a small MLP trained with Adam scores well above 92%.

Constraints and notes

• nn.CrossEntropyLoss takes raw logits, not probabilities.
• Predict with argmax under torch.no_grad().
• Return any callable that maps (n, d) to labels of shape (n,).