BUG Fix: Adding norm before generating the prediction

Mouradost · Aug 14, 2023 · e86f364 · e86f364
1 parent 2bdbb46
commit e86f364
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/src/trainers/base.py b/src/trainers/base.py
@@ -131,7 +131,9 @@ def get_predictions(
                     x.to(self.device),
                 )
                 prediction = self.scaler.reverse(
-                    self.model(x=x, adj=adj, adj_hat=sim, idx=idx).cpu().detach()
+                    self.model(x=self.scaler.norm(x), adj=adj, adj_hat=sim, idx=idx)
+                    .cpu()
+                    .detach()
                 )
                 all_grounds.append(y.detach())
                 all_predictions.append(prediction)

diff --git a/src/trainers/full.py b/src/trainers/full.py
@@ -187,7 +187,7 @@ def loss(
         if self.model_pred_single:
             y = y[:, -1:]
 
-        contrib = self.e_model(x=x, idx=idx)
+        contrib = self.e_model(x=self.scaler.norm(x), idx=idx)
         e_loss = self.loss_fn(contrib, sim, dim=(0, 1, 2))
         loss = self.loss_fn(
             pred=self.scaler.reverse(
@@ -218,7 +218,10 @@ def get_predictions(
                 )
                 prediction = self.scaler.reverse(
                     self.model(
-                        x=x, adj=adj, adj_hat=self.e_model(x=x, idx=idx), idx=idx
+                        x=self.scaler.norm(x),
+                        adj=adj,
+                        adj_hat=self.e_model(x=self.scaler.norm(x), idx=idx),
+                        idx=idx,
                     )
                     .cpu()
                     .detach()