fixed failing tests (#571)

ciemss · Apr 16, 2024 · ed2f432 · ed2f432
1 parent fbd8b23
commit ed2f432
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/pyciemss/interfaces.py b/pyciemss/interfaces.py
@@ -474,7 +474,7 @@ def sample(
         model = CompiledDynamics.load(model_path_or_json)
 
         logging_times = torch.arange(
-            start_time + logging_step_size, end_time, logging_step_size
+            start_time, end_time + logging_step_size, logging_step_size
         )
 
         # Check that num_samples is a positive integer

diff --git a/tests/fixtures.py b/tests/fixtures.py
@@ -247,7 +247,7 @@ def check_result_sizes(
         assert isinstance(v, torch.Tensor)
 
         num_timesteps = len(
-            torch.arange(start_time + logging_step_size, end_time, logging_step_size)
+            torch.arange(start_time, end_time + logging_step_size, logging_step_size)
         )
         if v.ndim == 2 and k == "model_weights":
             assert v.shape[0] == num_samples