Add some finess to the hyperparameter scan

Code/python/src/aiRNN/dataloader.py

@@ -126,11 +126,12 @@ class BaseDataset(Dataset):
         torch.save(data_dict, filepath)
 
 class SaveDataset(Dataset):
-    def __init__(self, data_dict, device="cpu"):
+    def __init__(self, data_dict, device="cpu", step=1):
         super().__init__()
         self.X_feat = data_dict["X_feat"].to(device)
         self.X_time = data_dict["X_time"].to(device)
         self.Y_out = data_dict["Y_out"].to(device)
+        self.step = step
         self.device = device
         if "X_context" in data_dict:
             self.X_context = data_dict["X_context"].to(device)
@@ -138,14 +139,14 @@ class SaveDataset(Dataset):
             self.X_context = None
 
     def __len__(self):
-        return self.X_feat.shape[0]
+        return self.X_feat.shape[0] // self.step
 
     def __getitem__(self, idx):
-        X_f = self.X_feat[idx]
-        X_t = self.X_time[idx]
-        Y = self.Y_out[idx]
+        X_f = self.X_feat[idx * self.step]
+        X_t = self.X_time[idx * self.step]
+        Y = self.Y_out[idx * self.step]
         if self.X_context is not None:
-            X_c = self.X_context[idx]
+            X_c = self.X_context[idx * self.step]
         else:
             X_c = None
         return X_f, X_t, Y, X_c

Code/python/src/aiRNN/models.py

@@ -9,9 +9,10 @@ class BaseRNN(nn.Module):
 
         base_unit = lambda in_size, out_size: nn.Sequential(
             nn.Linear(in_size, out_size, device=device),
-            nn.Tanh(),
+            nn.GELU(),
             nn.Linear(out_size, out_size, device=device),
-            nn.Tanh(),
+            nn.GELU(),
+            nn.LayerNorm(out_size, device=device)
         )
 
         self.time_proj = base_unit(time_in, hidden_size)
@@ -48,11 +49,7 @@ class BaseRNN(nn.Module):
             raise ValueError("Unsupported rnn_type")
 
         readout_in = rnn_size
-        self.readout = nn.Sequential(
-            nn.Linear(readout_in, hidden_size, device=device),
-            nn.Tanh(),
-            nn.Linear(hidden_size, out_size, device=device),
-        )
+        self.readout = nn.Linear(readout_in, out_size, device=device)
         self.context_in = context_in
         self.rnn_size = rnn_size