def forward(self, x):


    x=x.unsqueeze(1).float()
    out = self.layer1(x)
    
    out = self.layer2(out)
    
    out = out.view(out.size(0), -1)
    
    out = self.fc1(out)
    out=self.drop_out1(out)
    out=self.relu(out)

    out = self.fc2(out)
    out=self.drop_out2(out)
    out=self.relu(out)
    
    out = self.fc3(out)

    return out

and 

criterion=nn.BCEWithLogitsLoss()    indeed the input of the this loss function should be the out put from linear layer ?