fix: Fixed loss weight buffer (#136)

frgfm · Jul 28, 2021 · e1fed60 · e1fed60
1 parent afae658
commit e1fed60
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/holocron/nn/modules/loss.py b/holocron/nn/modules/loss.py
@@ -25,11 +25,13 @@ def __init__(
         # Cast class weights if possible
         self.weight: Optional[Tensor]
         if isinstance(weight, (float, int)):
-            self.weight = torch.Tensor([weight, 1 - weight])  # type: ignore[assignment]
+            self.register_buffer('weight', torch.Tensor([weight, 1 - weight]))
         elif isinstance(weight, list):
-            self.weight = torch.Tensor(weight)  # type: ignore[assignment]
+            self.register_buffer('weight', torch.Tensor(weight))
+        elif isinstance(weight, Tensor):
+            self.register_buffer('weight', weight)
         else:
-            self.weight = weight
+            self.weight = None
         self.ignore_index = ignore_index
         # Set the reduction method
         if reduction not in ['none', 'mean', 'sum']: