Skip to content

Commit a934a1a

Browse files
add local_sagemaker_session to horovod test (#652)
1 parent 3df9353 commit a934a1a

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

tests/integ/test_horovod.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,13 +55,14 @@ def test_horovod(sagemaker_session, instance_type, tmpdir):
5555
[1, 2],
5656
(2, 1),
5757
(2, 2)])
58-
def test_horovod_local_mode(instances, processes, tmpdir):
58+
def test_horovod_local_mode(sagemaker_local_session, instances, processes, tmpdir):
5959
output_path = 'file://%s' % tmpdir
6060

6161
estimator = TensorFlow(entry_point=os.path.join(horovod_dir, 'test_hvd_basic.py'),
6262
role='SageMakerRole',
6363
train_instance_count=2,
6464
train_instance_type='local',
65+
sagemaker_session=sagemaker_local_session,
6566
py_version=integ.PYTHON_VERSION,
6667
script_mode=True,
6768
output_path=output_path,

0 commit comments

Comments
 (0)