Compute accuracy for 4d torch tensors

I want to compute the classification accuracy, while my targets and predictions have the shape (batch_size, nb_classes, x_size_image, y_size_image). The accuracy is the number of correctly classified pixels divided by the total number of pixels. Is there an easy way to do that?

I am assuming that the target and prediction batches are arranged like this - that there are x_size_image*y_size_image pixels, and that each pixel corresponds to one of the n classes which are in a one hot representation. if not then correct me.
so basically just convert it into (n of pixels , no. classes) by doing

src = src.view(-1, src.shape[1])
trg = trg.view(-1, src.shape[1])

to get the most likely output use argmax or softmax as you need,

trg = trg.argmax(1) 
src = src.argmax(1)

and just sum the indices where these 2 are equal and divide by the length of either.