Extracting weights from SGD algorithm
So I am implementing SGD for a binary classification problem. There are 2 classes of points and I want to plot the decision boundary but I'm not sure how to extract the weights from the code to plot it.
Here is the code:
def train_model(train_dl, model):
    # define the optimization
    criterion = nn.BCELoss(reduction='none')
    optimizer = torch.optim.SGD(net.parameters(), lr=0.1)
    # enumerate epochs
    for epoch in range(10):
        # enumerate mini batches
        for i, (inputs, targets) in enumerate(train_dl):
            # clear the gradients
            optimizer.zero_grad()
            # compute the model output
            yhat = model(inputs)
            # calculate loss
            loss = criterion(yhat, targets)
            # credit assignment
            loss.backward()
            # update model weights
            optimizer.step()

Any help would be much appreciated! Thanks.

You do not extract this information from the SGD optimizer, this information is part of your model.
What you can do, at test time, is generate a grid of points, compute their prediction using the trained model and then plot the grid points coloring them according to the prediction.
