Just flatten everything in one order, let’s say your final feature map is 7 x 7, batch size is 4, class number is 80. Then the output tensor should be 4 x 80 x 7 x 7. Here is the step to compute the loss:
# Flatten the batch size and 7x7 feature map to one dimension
out = out.permute(0, 2, 3, 1).contiguous().view(-1, class_numer) # size is 196 x 80
# What about the order? Why you did permute? And what the *ell is contiguous?
# Quick answer is: The order of dimension one is as expected, which is
# First row of the 7x7 feature map of first batch
# Second row of feature map of first batch
# Thrid row of first batch
# ...
# First row of second batch
# ...
# Last row of last batch

# You can refer to https://github.com/pytorch/pytorch/issues/2013
# for how pytorch's `view` works, and why I put a permute and contiguous there


# Now let's regard the `out` as we 196 predictions 
# and for every prediction, we have 80 classes.
# And I believe that you can arrange the same order 
# targets from the ground truth, which should be
# a vector of 196 composed by real class numbers.
targets = tatgets_vector() # implement this function

# get loss, no need to do softmax
loss = F.cross_entropy(out, target)


I have re-implemented S3FD 6. It’s pretty like SSD, both are anchor(default box) based detector. Here is what I done for this https://github.com/louis-she/sfd.pytorch/blob/master/trainer.py#L105 26