diff --git a/src/data/smart_meter.py b/src/data/smart_meter.py
index 2b2cc16..37def80 100644
--- a/src/data/smart_meter.py
+++ b/src/data/smart_meter.py
@@ -23,23 +23,32 @@ def collate_fns(max_num_context, max_num_extra_target, sample, sort=True):
         x = torch.from_numpy(x).float()
         y = torch.from_numpy(y).float()
 
+        x[:,:max_num_context, -1] = 0  # Feature to let the model know this is past data
+        n=x[:, max_num_context:, -1].shape[1]
+        x[:, max_num_context:, -1] = torch.arange(1, n+1)/1.0/n  # Feature to let the model know this is past data
         x_context = x[:, :max_num_context]
         y_context = y[:, :max_num_context]
-
+    
         if sample:
             x_target_extra = x[:, max_num_context:]
             y_target_extra = y[:, max_num_context:]
 
-            # This is slightly differen't than normal, we are ensuring that out target point are in the future, to mimic deployment
+            # This is slightly differen't than normal, we are ensuring that our target point are in the future, to mimic deployment
             x_context, y_context = npsample_batch(
                 x_context, y_context, size=num_context, sort=sort
             )
+
             x_target_extra, y_target_extra = npsample_batch(
                 x_target_extra, y_target_extra, size=num_extra_target, sort=sort
             )
 
             x = torch.cat([x_context, x_target_extra], 1)
-            y = torch.cat([y_context, y_target_extra], 1)
+            y = torch.cat([y_context, y_target_extra], 1)            
+
+        assert (x_context[:, :, -1]==0).all()
+        assert (x[:, -1, -1] > 0).all()
+        assert (x[:, 0, -1] == 0).all()
+        
         return x_context, y_context, x, y
 
     return collate_fn
@@ -61,6 +70,9 @@ class SmartMeterDataSet(torch.utils.data.Dataset):
         columns = ['tstp'] + list(set(rows.columns) - set(['tstp']))
         rows = rows[columns]
 
+        # This will be the last row, and will change it upon sample to let the model know some points are in the future
+        rows['future']=1
+
         x = rows.drop(columns=self.label_names)
         y = rows[self.label_names]
         return x, y
diff --git a/src/models/lstm.py b/src/models/lstm.py
index c8d052c..0346744 100644
--- a/src/models/lstm.py
+++ b/src/models/lstm.py
@@ -2,6 +2,7 @@ import os
 import numpy as np
 import pandas as pd
 import torch
+from tqdm.auto import tqdm
 from torch import nn
 from torch.nn import functional as F
 from torch.utils.data import DataLoader
@@ -140,6 +141,13 @@ class LSTM_PL(pl.LightningModule):
         assert torch.isfinite(avg_loss)
         return {"avg_val_loss": avg_loss, "log": tensorboard_logs}
 
+
+    def test_step(self, *args, **kwargs):
+        return self.validation_step(*args, **kwargs)
+
+    def test_end(self, *args, **kwargs):
+        return self.validation_end(*args, **kwargs)
+
     def configure_optimizers(self):
         optim = torch.optim.Adam(self.parameters(), lr=self.hparams["learning_rate"])
         scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
@@ -219,13 +227,13 @@ class LSTM_PL(pl.LightningModule):
         return parser
 
 
-def plot_from_loader(loader, model, vis_i=670):
+def plot_from_loader(loader, model, vis_i=670, n=100):
     dset_test = loader.dataset
     label_names = dset_test.label_names
     y_trues = []
     y_preds = []
     vis_i = min(vis_i, len(dset_test))
-    for i in range(vis_i, vis_i + 100):
+    for i in tqdm(range(vis_i, vis_i + n)):
         x_rows, y_rows = dset_test.iloc(i)
         x, y = dset_test[i]
         device = next(model.parameters()).device