Fix bug in steps for dataset

2025-11-23 20:59:42 +01:00
parent 30e21eb5e9
commit b012a86768
1 changed files with 5 additions and 5 deletions
@@ -139,14 +139,14 @@ class SaveDataset(Dataset):
            self.X_context = None

    def __len__(self):
-        return self.X_feat.shape[0] // self.step
+        return self.X_feat.shape[0]

    def __getitem__(self, idx):
-        X_f = self.X_feat[idx * self.step]
-        X_t = self.X_time[idx * self.step]
-        Y = self.Y_out[idx * self.step]
+        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 * self.step]
+            X_c = self.X_context[idx]
        else:
            X_c = None
        return X_f, X_t, Y, X_c