To create the predictions containing the class index associated with the highest probability (or logit) you could use:
preds = torch.argmax(output, dim=0) # assuming output has the mentioned shape [nb_classes, height, weight]