update unittest fixtures

shaoanlu · Apr 28, 2024 · 934d610 · 934d610
1 parent 6fccd8e
commit 934d610
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/tests/config/test_config.yaml b/tests/config/test_config.yaml
@@ -39,7 +39,7 @@ trainer:
         num_warmup_steps: 500
 
 dataloader:
-    batch_size: 8
+    batch_size: 3
 
 normalizer:
     action:

diff --git a/tests/test_dataset.joblib b/tests/test_dataset.joblib
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
@@ -24,20 +24,21 @@ def test_init(self):
         self.assertTrue(isinstance(dataset, PlanarQuadrotorStateDataset))
 
     def test_iter(self):
+        batch_size = self.config["dataloader"]["batch_size"]
         dataset = PlanarQuadrotorStateDataset(dataset_path=self.dataset_path, config=self.config)
         dataloader = torch.utils.data.DataLoader(
             dataset,
-            batch_size=self.config["dataloader"]["batch_size"],
+            batch_size=batch_size,
             shuffle=True,
             pin_memory=True,
         )
 
         # batch context matches expectecd shapes
         batch = next(iter(dataloader))
         self.assertEqual(
-            batch["obs"].shape, (self.batch_size, self.obs_dim * self.obs_horizon + self.obstacle_encode_dim)
+            batch["obs"].shape, (batch_size, self.obs_dim * self.obs_horizon + self.obstacle_encode_dim)
         )
-        self.assertEqual(batch["action"].shape, (self.batch_size, self.pred_horizon, self.action_dim))
+        self.assertEqual(batch["action"].shape, (batch_size, self.pred_horizon, self.action_dim))
 
 
 if __name__ == "__main__":