# Codeblock 5e
    target_class = target[..., :20]      #(1)
    pred_class = prediction[..., :20]    #(2)
    
    
    class_loss = sse(      #(3)
        torch.flatten(obj * target_class, end_dim=-2),
        torch.flatten(obj * pred_class, end_dim=-2),
    )