SUBTASK 4 Training Integration

[shanjiaz]

Originally a task in Add P-eagle support in training.

Task

• Add cross entropy loss function as an option in training script
• Group-specific hidden state and target extraction at sampled indices
• Per-depth accuracy metrics
• Training loop modifications for single-pass parallel group processing