帮我用pytorch定义一个attention的函数

好的，以下是一个简单的 PyTorch Attention 函数的定义：

import torch
import torch.nn.functional as F

def attention(query, key, value):
    scores = torch.matmul(query, key.transpose(-2, -1))
    scores = scores / torch.sqrt(torch.tensor(query.size(-1)).float())
    weights = F.softmax(scores, dim=-1)
    output = torch.matmul(weights, value)
    return output

这个函数接受三个输入参数：query、key 和 value。它计算出一个注意力权重矩阵，然后将该矩阵与 value 矩阵相乘，得到最终的输出。