如何在PyTorch中计算注意力得分和编码器输出的加权平均值？

import torch
import torch.nn as nn

# 定义注意力机制模型
class Attention(nn.Module):
    def __init__(self, hidden_size):
        super(Attention, self).__init__()
        self.hidden_size = hidden_size
        self.query = nn.Linear(hidden_size, hidden_size)
        self.key = nn.Linear(hidden_size, hidden_size)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, query, key):
        q = self.query(query)
        k = self.key(key)
        attention_scores = torch.matmul(q, k.transpose(1, 2))
        attention_weights = self.softmax(attention_scores)
        return attention_weights

# 定义编码器模型
class Encoder(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(Encoder, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.embedding = nn.Embedding(input_size, hidden_size)
        self.gru = nn.GRU(hidden_size, hidden_size)

    def forward(self, input):
        embedded = self.embedding(input)
        output, hidden = self.gru(embedded)
        return output, hidden

# 定义注意力加权平均模型
class AttentionWeightedAverage(nn.Module):
    def __init__(self, hidden_size):
        super(AttentionWeightedAverage, self).__init__()
        self.hidden_size = hidden_size
        self.attention = Attention(hidden_size)

    def forward(self, query, key):
        attention_weights = self.attention(query, key)
        weighted_average = torch.matmul(attention_weights, key)
        return weighted_average

# 示例数据
input_size = 100
hidden_size = 256
batch_size = 32
seq_length = 10

# 创建编码器实例
encoder = Encoder(input_size, hidden_size)

# 创建注意力加权平均实例
attention_weighted_average = AttentionWeightedAverage(hidden_size)

# 随机生成输入数据
input = torch.randint(0, input_size, (batch_size, seq_length))

# 计算编码器输出
output, hidden = encoder(input)

# 计算注意力加权平均值
weighted_average = attention_weighted_average(hidden, output)

print(weighted_average)