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


class SoftCrossEntropyLoss(nn.Module):
    """Computes the CrossEntropyLoss while accepting soft (float) targets

    Args:
        weight: a tensor of relative weights to assign to each class.
        reduction: how to combine the elmentwise losses
            'none': return an unreduced list of elementwise losses
            'elementwise_mean': return the mean loss per elements
            'sum': return the sum of the elementwise losses

    Accepts:
        input: An [n, k] float tensor of prediction logits (not probabilities)
        target: An [n, k] float tensor of target probabilities
    """

    def __init__(self, weight=None, reduction="elementwise_mean"):
        super().__init__()
        assert weight is None or isinstance(weight, torch.FloatTensor)
        self.weight = weight
        if self.weight is not None:
            self.weight = torch.Tensor(self.weight)
        self.reduction = reduction

    def forward(self, input, target):
        n, k = input.shape
        # Note that t.new_zeros, t.new_full put tensor on same device as t
        cum_losses = input.new_zeros(n)
        # Moving weights to appropriate device
        if self.weight is not None:
            self.weight = self.weight.type(input.dtype).to(input.device)
        for y in range(k):
            cls_idx = input.new_full((n,), y, dtype=torch.long)
            y_loss = F.cross_entropy(input, cls_idx, reduction="none")
            if self.weight is not None:
                y_loss = y_loss * self.weight[y]
            cum_losses += target[:, y].float() * y_loss
        if self.reduction == "none":
            return cum_losses
        elif self.reduction == "elementwise_mean":
            return cum_losses.mean()
        elif self.reduction == "sum":
            return cum_losses.sum()
        else:
            raise ValueError(f"Unrecognized reduction: {self.reduction}")

@@ -768,79 +768,8 @@
 _()%0A
-        assert weight is None or isinstance(weight, torch.FloatTensor)%0A
     
@@ -881,16 +881,97 @@
 weight)%0A
+        assert self.weight is None or isinstance(self.weight, torch.FloatTensor)%0A
         

Fix faulty assert statement