You’re getting that error in the line where you calculate the loss.
You’re using the BCEWithLogitsLoss(). The input and output have to be the same size and have the dtype float.
y_pred = (batch_size, *) , Float
y_train = (batch_size, *) , Float
What you need to do it, return only a single output value from your network. To do that, you need to change this line self.fc2 = nn.Linear(512, 2) to self.fc2 = nn.Linear(512, 1).
Once you do this, your network will return an output of size [10, 1]. Now, you will still get an error because your target is of the size [10]. So, run .squeeze() on the predicted tensor which will remove the dimension ‘1’ and convert it to [10] from [10,1] ie. in your train loop, do the following before you calculate loss -
    for epoch in range(EPOCHS):        
        for step, (x, y) in enumerate(train_data_loader):
            b_x = Variable(x)   # batch x (image)
            b_y = Variable(y)   # batch y (target)
            output = model(b_x)[0]          
            loss = loss_func(output.squeeze(), b_y)
            ...

Let me know if this works.