32

batch_size = prediction.size(0)
write = False
    
 # we can do non max suppression only on individual images so we will loop through images
 for ind in range(batch_size):  
     image_pred = prediction[ind] 
     # we will take only those rows with maximm class probability
     # and corresponding index
     max_conf, max_conf_score = torch.max(image_pred[:,5:5+ num_classes], 1)
     max_conf = max_conf.float().unsqueeze(1)
     max_conf_score = max_conf_score.float().unsqueeze(1)
     combined = (image_pred[:,:5], max_conf, max_conf_score)
     # concatinating index values and max probability with box cordinates as columns
     image_pred = torch.cat(combined, 1)