Add sync_dist to SSL Models

Author: sorenmacbeth

src/pytorch_tabular/ssl_models/base_model.py

@@ -136,11 +136,11 @@ def _setup_metrics(self):
         pass
 
     @abstractmethod
-    def calculate_loss(self, output, tag):
+    def calculate_loss(self, output, tag, sync_dist):
         pass
 
     @abstractmethod
-    def calculate_metrics(self, output, tag):
+    def calculate_metrics(self, output, tag, sync_dist):
         pass
 
     @abstractmethod
@@ -167,15 +167,15 @@ def training_step(self, batch, batch_idx):
     def validation_step(self, batch, batch_idx):
         with torch.no_grad():
             output = self.forward(batch)
-            self.calculate_loss(output, tag="valid")
-            self.calculate_metrics(output, tag="valid")
+            self.calculate_loss(output, tag="valid", sync_dist=True)
+            self.calculate_metrics(output, tag="valid", sync_dist=True)
         return output
 
     def test_step(self, batch, batch_idx):
         with torch.no_grad():
             output = self.forward(batch)
-            self.calculate_loss(output, tag="test")
-            self.calculate_metrics(output, tag="test")
+            self.calculate_loss(output, tag="test", sync_dist=True)
+            self.calculate_metrics(output, tag="test", sync_dist=True)
         return output
 
     def on_validation_epoch_end(self) -> None:

src/pytorch_tabular/ssl_models/dae/dae.py

@@ -200,7 +200,7 @@ def forward(self, x: Dict):
             else:
                 return z.features
 
-    def calculate_loss(self, output, tag):
+    def calculate_loss(self, output, tag, sync_dist=False):
         total_loss = 0
         for type_, out in output.items():
             if type_ == "categorical":
@@ -220,6 +220,7 @@ def calculate_loss(self, output, tag):
                 on_step=False,
                 logger=True,
                 prog_bar=False,
+                sync_dist=sync_dist,
             )
             total_loss += loss
         self.log(
@@ -230,10 +231,11 @@ def calculate_loss(self, output, tag):
             # on_step=False,
             logger=True,
             prog_bar=True,
+            sync_dist=sync_dist,
         )
         return total_loss
 
-    def calculate_metrics(self, output, tag):
+    def calculate_metrics(self, output, tag, sync_dist=False):
         pass
 
     def featurize(self, x: Dict):