fix: return latent vars if not in the training stage;

Author: WenjieDu

pypots/imputation/saits/model.py

@@ -347,21 +347,18 @@ def predict(
                 results = self.model.forward(
                     inputs, diagonal_attention_mask, training=False
                 )
-                imputed_data = results["imputed_data"]
-                imputation_collector.append(imputed_data)
+                imputation_collector.append(results["imputed_data"])
 
                 if return_latent_vars:
-                    first_DMSA_attn_weights = (
+                    first_DMSA_attn_weights_collector.append(
                         results["first_DMSA_attn_weights"].cpu().numpy()
                     )
-                    second_DMSA_attn_weights = (
+                    second_DMSA_attn_weights_collector.append(
                         results["second_DMSA_attn_weights"].cpu().numpy()
                     )
-                    combining_weights = results["combining_weights"].cpu().numpy()
-
-                    first_DMSA_attn_weights_collector.append(first_DMSA_attn_weights)
-                    second_DMSA_attn_weights_collector.append(second_DMSA_attn_weights)
-                    combining_weights_collector.append(combining_weights)
+                    combining_weights_collector.append(
+                        results["combining_weights"].cpu().numpy()
+                    )
 
         # Step 3: output collection and return
         imputation = torch.cat(imputation_collector).cpu().detach().numpy()

pypots/imputation/saits/modules/core.py

@@ -174,11 +174,15 @@ def forward(
             [first_DMSA_attn_weights, second_DMSA_attn_weights, combining_weights],
         ) = self._process(inputs, diagonal_attention_mask)
 
+        results = {
+            "first_DMSA_attn_weights": first_DMSA_attn_weights,
+            "second_DMSA_attn_weights": second_DMSA_attn_weights,
+            "combining_weights": combining_weights,
+            "imputed_data": imputed_data,
+        }
         if not training:
             # if not in training mode, return the classification result only
-            return {
-                "imputed_data": imputed_data,
-            }
+            return results
 
         ORT_loss = 0
         ORT_loss += self.customized_loss_func(X_tilde_1, X, masks)
@@ -193,13 +197,10 @@ def forward(
         # `loss` is always the item for backward propagating to update the model
         loss = self.ORT_weight * ORT_loss + self.MIT_weight * MIT_loss
 
-        results = {
-            "first_DMSA_attn_weights": first_DMSA_attn_weights,
-            "second_DMSA_attn_weights": second_DMSA_attn_weights,
-            "combining_weights": combining_weights,
-            "imputed_data": imputed_data,
-            "ORT_loss": ORT_loss,
-            "MIT_loss": MIT_loss,
-            "loss": loss,  # will be used for backward propagating to update the model
-        }
+        results["ORT_loss"] = ORT_loss
+        results["MIT_loss"] = MIT_loss
+
+        # will be used for backward propagating to update the model
+        results["loss"] = loss
+
         return results