-- moved nested function out for pickling

Author: manujosephv

pytorch_tabular/models/mixture_density/mdn.py

@@ -378,6 +378,9 @@ def unpack_input(self, x: Dict):
 class NODEMDN(BaseMDN):
     def __init__(self, config: DictConfig, **kwargs):
         super().__init__(config, **kwargs)
+    
+    def subset(self, x):
+        return x[..., :].mean(dim=-2)
 
     def _build_network(self):
         self.hparams.node_input_dim = (
@@ -387,10 +390,7 @@ def _build_network(self):
         # average first n channels of every tree, where n is the number of output targets for regression
         # and number of classes for classification
 
-        def subset(x):
-            return x[..., :].mean(dim=-2)
-
-        output_response = utils.Lambda(subset)
+        output_response = utils.Lambda(self.subset)
         self.backbone = nn.Sequential(backbone, output_response)
         # Adding the last layer
         self.hparams.mdn_config.input_dim = backbone.output_dim

pytorch_tabular/models/node/node_model.py

@@ -60,6 +60,9 @@ def __init__(self, config: DictConfig, **kwargs):
         if config.embed_categorical:
             self.embedding_cat_dim = sum([y for x, y in config.embedding_dims])
         super().__init__(config, **kwargs)
+    
+    def subset(self, x):
+            return x[..., : self.hparams.output_dim].mean(dim=-2)
 
     def _build_network(self):
         if self.hparams.embed_categorical:
@@ -79,10 +82,7 @@ def _build_network(self):
         # average first n channels of every tree, where n is the number of output targets for regression
         # and number of classes for classification
 
-        def subset(x):
-            return x[..., : self.hparams.output_dim].mean(dim=-2)
-
-        self.output_response = utils.Lambda(subset)
+        self.output_response = utils.Lambda(self.subset)
 
     def unpack_input(self, x: Dict):
         if self.hparams.embed_categorical: