在pytorch中屏蔽和计算通过具有线性输出层的RNN发送的填充批次的损失

import torch
import torch.nn as nn
import torch.nn.utils.rnn as rnn_utils

# 假设输入序列为input_seq，标签为targets，长度为seq_lengths
input_seq = torch.tensor([[1, 2, 3], [4, 5, 0], [6, 0, 0]])
targets = torch.tensor([1, 2, 3])
seq_lengths = torch.tensor([3, 2, 1])

# 定义RNN模型
class RNNModel(nn.Module):
    def __init__(self):
        super(RNNModel, self).__init__()
        self.rnn = nn.RNN(input_size=3, hidden_size=5, batch_first=True)
        self.linear = nn.Linear(5, 4)
    
    def forward(self, input_seq, seq_lengths):
        packed_seq = rnn_utils.pack_padded_sequence(input_seq, seq_lengths, batch_first=True, enforce_sorted=False)
        rnn_output, _ = self.rnn(packed_seq)
        unpacked_output, _ = rnn_utils.pad_packed_sequence(rnn_output, batch_first=True)
        linear_output = self.linear(unpacked_output)
        return linear_output

# 创建模型实例
model = RNNModel()

# 计算输出
output = model(input_seq, seq_lengths)

# 根据原始索引顺序对输出进行排序
sorted_output = output[torch.argsort(seq_lengths, descending=True)]

# 计算损失函数
loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(sorted_output, targets)

# 打印损失
print(loss)