如何在PyTorch中计算自举交叉熵损失？

import torch
import torch.nn as nn
import torch.nn.functional as F

class BootstrapCrossEntropyLoss(nn.Module):
    def __init__(self, num_classes, num_bootstrap_samples, alpha):
        super(BootstrapCrossEntropyLoss, self).__init__()
        self.num_classes = num_classes
        self.num_bootstrap_samples = num_bootstrap_samples
        self.alpha = alpha

    def forward(self, inputs, targets):
        batch_size = inputs.size(0)
        bootstrap_targets = targets.repeat(self.num_bootstrap_samples)
        bootstrap_inputs = inputs.repeat(self.num_bootstrap_samples, 1)

        log_probs = F.log_softmax(bootstrap_inputs, dim=1)
        probs = torch.exp(log_probs)

        bootstrap_loss = F.nll_loss(log_probs, bootstrap_targets, reduction='none')
        bootstrap_loss = bootstrap_loss.view(self.num_bootstrap_samples, batch_size)
        bootstrap_loss = torch.mean(bootstrap_loss, dim=0)

        weights = torch.zeros_like(targets, dtype=torch.float)
        for i in range(self.num_classes):
            class_mask = targets == i
            class_samples = torch.sum(class_mask).item()
            class_weight = (1 - self.alpha) / class_samples + self.alpha / self.num_classes
            weights += class_mask.float() * class_weight

        weighted_loss = torch.mean(weights * bootstrap_loss)
        return weighted_loss

model = YourModel()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

criterion = BootstrapCrossEntropyLoss(num_classes, num_bootstrap_samples, alpha)

for epoch in range(num_epochs):
    # 前向传播
    outputs = model(inputs)
    loss = criterion(outputs, targets)

    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()