transformers - 💡(How to fix) Fix Investigate training support for Sapiens2ForPoseEstimation when labels are provided

Feature request

While exploring the recently added Sapiens2ForPoseEstimation implementation, I noticed that the model's forward() method accepts labels (documented as heatmap ground truth with shape (batch_size, num_keypoints, height, width)), but currently raises NotImplementedError("Training is not yet supported") when labels are provided.

I investigated the original Meta Sapiens2 pose-estimation implementation and found that pose training is configured through OpenMMLab's KeypointMSELoss with use_target_weight=True.

Relevant references:

• Hugging Face implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/models/sapiens2/modular_sapiens2.py#L1713

• Meta Sapiens2 pose configuration: https://github.com/facebookresearch/sapiens2/blob/main/sapiens/pose/configs/keypoints308/shutterstock_goliath_3po/sapiens2_1b_keypoints308_shutterstock_goliath_3po-1024x768.py#L270-297

• OpenMMLab KeypointMSELoss implementation: https://github.com/open-mmlab/mmpose/blob/main/mmpose/models/losses/heatmap_loss.py

• Sapiens2 paper: https://arxiv.org/abs/2604.21681

I would like to better understand whether the lack of training support is intentional or if support for supervised pose-estimation training would be considered within scope for the Transformers implementation.

Motivation

While exploring Sapiens2ForPoseEstimation, I noticed that the model accepts labels and documents them as heatmap ground truth targets for loss computation, but currently raises NotImplementedError("Training is not yet supported") whenever labels are provided.

To understand whether this was specific to Sapiens2 or part of a broader design decision, I investigated the original Meta Sapiens2 implementation as well as the OpenMMLab components used during training.

My findings were:

• The Hugging Face implementation already supports the complete inference pipeline, including feature extraction, heatmap decoding, and returning pose heatmaps.
• The forward signature explicitly accepts labels and documents them as heatmap supervision targets.
• The original Meta Sapiens2 pose-estimation configuration uses OpenMMLab's KeypointMSELoss with use_target_weight=True and loss_weight=10.0.
• The corresponding OpenMMLab implementation computes a heatmap-based MSE loss and supports target weighting through a masking mechanism used for keypoint visibility.
• I also checked VitPoseForPoseEstimation in Transformers and found the same pattern: labels are accepted but training currently raises NotImplementedError.

Because both Sapiens2 and ViTPose expose pose-estimation outputs and accept supervision targets but currently stop before loss computation, I wanted to understand whether pose-estimation training support is intentionally omitted from Transformers or whether support for supervised training would be considered within scope.

Your contribution

I would be happy to contribute a PR if maintainers consider pose-estimation training support to be within scope for Transformers.

Based on the investigation above, a potential direction would be to implement the supervised pose-estimation loss path when labels are provided, using the original Sapiens2 training setup and OpenMMLab's KeypointMSELoss behavior as references while remaining consistent with existing Transformers conventions.

If this aligns with maintainer expectations, I would be interested in contributing the implementation, accompanying tests, and documentation.