How to use pytorch to output the probability of binary classfication?

@ptrblck
Hi,thanks!
I use this code:

import torch
import torch.nn as nn
import torch.nn.functional as F


train = pd.read_csv(r'train.csv')
x_test = pd.read_csv(r'test.csv')

x_train = train[:, :-1]
y_train = train[:, -1]

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(800000, 153)
        self.fc2 = nn.Linear(153, 2)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.BCEWithLogitsLoss()
x_train = torch.Tensor(x_train.values)
y_train = torch.Tensor(y_train.values)
x_test = torch.Tensor(x_test.values)
# data = torch.randn(64, 128)
# target = torch.randint(0, 2, (64, 40)).float()

for epoch in range(5):
    optimizer.zero_grad()
    output = model(x_train)
    loss = criterion(output, y_train)
    loss.backward()
    optimizer.step()
    print('epoch {}, loss {}'.format(epoch, loss.item()))

# check predictions
output = model(x_test)
probs = torch.sigmoid(output)
print(probs)

But the error messgae is:

RuntimeError:
size mismatch, m1: [800000 x 153], m2: [800000 x 153] at …\aten\src\TH/generic/THTensorMath.cpp:41

How to set these 2 lines?

self.fc1 = nn.Linear(800000, 153)
self.fc2 = nn.Linear(153, 2)