我想要用pytorch实现一个加入attention机制的RNN，请帮我完善以下代码：

当然，我会尽力帮助您。以下是您可能会用到的代码示例：

import torch
import torch.nn as nn

class AttentionRNN(nn.Module):

def __init__(self, input_size, hidden_size, num_layers, dropout=0.5):
super(AttentionRNN, self).__init__()

self.hidden_size = hidden_size
self.num_layers = num_layers

self.embedding = nn.Embedding(input_size, hidden_size)
self.gru = nn.GRU(hidden_size, hidden_size, num_layers, batch_first=True, dropout=dropout)

# attention mechanism
self.attention = nn.Linear(hidden_size * 2, 1)
self.softmax = nn.Softmax(dim=1)
self.dropout = nn.Dropout(dropout)
self.linear = nn.Linear(hidden_size, hidden_size)

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

# attention mechanism
batch_size = output.size(0)
output = output.view(-1, self.hidden_size) # flatten
hidden = hidden[-1] # last layer's hidden state
hidden = hidden.unsqueeze(1).repeat(1, output.size(1), 1) # repeat hidden state for all timesteps
cat = torch.cat((output, hidden), dim=2) # concatenate output and hidden
energy = self.attention(cat).squeeze(2) # calculate energy for each timestep
attention_scores = self.softmax(energy) # apply softmax to get attention weights
attention_scores = attention_scores.unsqueeze(2)
context = torch.bmm(output.permute(0,2,1), attention_scores).squeeze(2) # calculate context vector
output = self.linear(self.dropout(context)) # apply linear transformation

return output, hidden

def init_hidden(self, batch_size):
return torch.zeros(self.num_layers, batch_size, self.hidden_size)