Question in pytorch

how I create the final result after the SOFTMAX function of the CNN network in binary image form (0,1)?. With the entry of this network is a single image and from this image we reserve (80% train, 20% test).

Help me plzzzz

You can use torch.argmax and either pass the logits or the probabilities to it in case you are working on a multi-class classification.

Yes how ? plzz @ptrblck

I want to display the output image

Assuming your model outputs logits in the shape [batch_size, nb_classes ,*], where * denotes additional dimensions, this would work:

output = model(input)
pred = torch.argmax(output, dim=1)