Self-Attention Weighted Aggregation - Coding problem | OneNoughtOne