🐛Fixed legacy bug (dc57b0f)

carefree0910 · Jan 18, 2021 · 745cd42 · 745cd42
1 parent 9a1bf14
commit 745cd42
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/cflearn/modules/heads/traditional/core.py b/cflearn/modules/heads/traditional/core.py
@@ -118,8 +118,9 @@ def __init__(
         y_ravel: np.ndarray,
         categorical: Optional[torch.Tensor],
         categorical_dims: Dict[int, int],
+        **kwargs: Any,
     ):
-        super().__init__(in_dim, out_dim)
+        super().__init__(in_dim, out_dim, **kwargs)
         self.categorical_dims = categorical_dims
         if categorical is None:
             self.mnb = None
@@ -195,8 +196,9 @@ def __init__(
         y_ravel: np.ndarray,
         pretrain: bool,
         numerical: Optional[torch.Tensor],
+        **kwargs: Any,
     ):
-        super().__init__(in_dim, out_dim)
+        super().__init__(in_dim, out_dim, **kwargs)
         if numerical is None:
             self.mu = self.std = self.normal = None
         else: