Marco Rossi. Multi-Head Attention-Based OOM Risk Prediction for CNN and Transformer Training Tasks on GPUMemNet. JAIR. 2026;7(2):148-163. doi:10.69987/