fix: add super().__init__ in the ContrastiveLoss Modules

ruflab · Jan 23, 2024 · 41f8f6d · 41f8f6d
1 parent 6f9f6c5
commit 41f8f6d
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/shimmer/modules/losses.py b/shimmer/modules/losses.py
@@ -51,6 +51,8 @@ def __init__(
         logit_scale: torch.Tensor,
         reduction: Literal["mean", "sum", "none"] = "mean",
     ) -> None:
+        super().__init__()
+
         self.register_buffer("logit_scale", logit_scale)
         self.reduction: Literal["mean", "sum", "none"] = reduction
 
@@ -88,6 +90,8 @@ def __init__(
         logit_scale: torch.Tensor,
         reduction: Literal["mean", "sum", "none"] = "mean",
     ) -> None:
+        super().__init__()
+
         self.register_buffer("logit_scale", logit_scale)
         self.reduction: Literal["mean", "sum", "none"] = reduction