Model Configuration, Training, and Predictions for Classification

Model configuration

In the chapter, Going Classy, we ended up with a lean model configuration part. We only need to define a model, an appropriate loss function, and an optimizer. Let us define a model that produces logits and uses BCEWithLogitsLoss as the loss function. Since we have two features, and we are producing logits instead of probabilities, our model has one layer and one layer alone: Linear(2, 1). We will keep using the SGD optimizer with a learning rate of 0.1 for now.

This is what the model configuration looks like for our classification problem: