Hi Raouf!
I don’t follow in detail what you are doing, but be aware that
the output of a model with a final sigmoid() layer will, of
course, be different than the output of the analogous model
that lacks that layer.
With the sigmoid() layer your model returns probabilities.
Without the sigmoid() layer, your model returns “raw scores”
that are called logits.
Some comments appear in line, below.



 raouf_ks:

I have a problem regarding “BCE With Logits Loss”
My simple neural net is
        x = self.output(x)
        x = torch.sigmoid(x)
        return x
...



Removing the x = torch.sigmoid(x) will, of course, change
the value x that your model returns.


criterion = nn.BCELoss()



The loss you calculate will be (more or less) the same if you
then use criterion = nn.BCEWithLogitsLoss() without
the sigmoid() layer, but the model’s output is different.  This
difference – logits instead of probabilities – is compensated for
by using the different loss function.


my training loop
...
            outputs = bnn(X.float())
            loss = criterion(outputs.float().reshape(-1), Y.float())
...
        plot_decision_boundary( bnn, X, Y, epoch, ((outputs.squeeze()>=0.5) == Y).float().mean(), 
                                nbh=4, model_type='mcdropout')



Using outputs.squeeze()>=0.5 as your binary predictions
is appropriate when using the final sigmoid() layer.  But unless
you change the logic for your prediction / accuracy calculation
when you remove the sigmoid() layer you will get different
(and incorrect) results.
When outputs are the logits you get without the sigmoid()
layer, you should use outputs.squeeze()>=0.0 to get your
binary predictions.


My plot function
def plot_decision_boundary(model, X, Y, epoch, accuracy, model_type='classic', samples=100, nbh=2, cmap='RdBu'):    
...
    if model_type=='classic':
        model.eval()
        pred = model(test_tensor)
...
    Z = pred.reshape(xx.shape).detach().numpy()

    plt.cla()
    ax.set_title('Classification Analysis')
    ax.contourf(xx, yy, Z, cmap=cmap, alpha=0.25)
    ax.contour(xx, yy, Z, colors='k', linestyles=':', linewidths=0.7)

My problem is if I keep the sigmoid in the forward and use BCELoss I get this final result
…
However if I change it without a sigmoid function and use BCEwithLogits I get this result
…


You plot Z.  Z comes from preds, which comes from your model.
When you remove the sigmoid() from your model, you change
preds and Z, so your plots come out different because you are
plotting something different.
Best.
K. Frank