trainboundary->_trainboundary

determined-ai · Oct 25, 2024 · 34fda27 · 34fda27
1 parent 1a7de27
commit 34fda27
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/harness/tests/experiment/pytorch/test_pytorch_trial.py b/harness/tests/experiment/pytorch/test_pytorch_trial.py
@@ -82,8 +82,8 @@ def test_onevar_single(self, tmp_path: pathlib.Path) -> None:
         train_steps, metrics = trial_controller._train_with_boundaries(
             training_enumerator=enumerate(trial_controller.training_iterator),
             train_boundaries=[
-                pytorch.TrainBoundary(
-                    step_type=pytorch.TrainBoundaryType.TRAIN, unit=pytorch.Batch(100)
+                pytorch._TrainBoundary(
+                    step_type=pytorch._TrainBoundaryType.TRAIN, unit=pytorch.Batch(100)
                 )
             ],
         )
@@ -114,8 +114,8 @@ def test_training_metrics(self, tmp_path: pathlib.Path) -> None:
         train_steps, metrics = trial_controller._train_with_boundaries(
             training_enumerator=enumerate(trial_controller.training_iterator),
             train_boundaries=[
-                pytorch.TrainBoundary(
-                    step_type=pytorch.TrainBoundaryType.TRAIN, unit=pytorch.Batch(100)
+                pytorch._TrainBoundary(
+                    step_type=pytorch._TrainBoundaryType.TRAIN, unit=pytorch.Batch(100)
                 )
             ],
         )
@@ -542,8 +542,8 @@ def test_variable_workload_size(self, tmp_path: pathlib.Path) -> None:
             train_steps, metrics = controller._train_with_boundaries(
                 training_enumerator=enumerate(controller.training_iterator),
                 train_boundaries=[
-                    pytorch.TrainBoundary(
-                        step_type=pytorch.TrainBoundaryType.TRAIN, unit=pytorch.Batch(num_batches)
+                    pytorch._TrainBoundary(
+                        step_type=pytorch._TrainBoundaryType.TRAIN, unit=pytorch.Batch(num_batches)
                     )
                 ],
             )