query = inputs[1]
attention_scores = inputs @ query
attention_weights = torch.softmax(attention_scores, dim=0)
context_vector = attention_weights @ inputs

        __
        __