diff --git a/metal/end_model/loss.py b/metal/end_model/loss.py
index <HASH>..<HASH> 100644
--- a/metal/end_model/loss.py
+++ b/metal/end_model/loss.py
@@ -20,10 +20,9 @@ class SoftCrossEntropyLoss(nn.Module):
 
     def __init__(self, weight=None, reduction="elementwise_mean"):
         super().__init__()
-        self.weight = weight
-        if self.weight is not None:
-            self.weight = torch.Tensor(self.weight)
-        assert self.weight is None or isinstance(self.weight, torch.FloatTensor)
+        self.weight = None
+        if weight is not None:
+            self.weight = torch.FloatTensor(weight)
         self.reduction = reduction
 
     def forward(self, input, target):
Remove unnecessary assert and convert weight to FloatTensor