Python PyTorchModel.prepare_container_defの例

コード例 #1

def test_model_image_accelerator(sagemaker_session):
    model = PyTorchModel(MODEL_DATA,
                         role=ROLE,
                         entry_point=SCRIPT_PATH,
                         sagemaker_session=sagemaker_session)
    with pytest.raises(ValueError):
        model.prepare_container_def(INSTANCE_TYPE,
                                    accelerator_type=ACCELERATOR_TYPE)

コード例 #2

ファイル: test_pytorch.py プロジェクト: ngluna/sagemaker-python-sdk

def test_model_prepare_container_def_no_instance_type_or_image():
    model = PyTorchModel(MODEL_DATA, role=ROLE, entry_point=SCRIPT_PATH)

    with pytest.raises(ValueError) as e:
        model.prepare_container_def()

    expected_msg = "Must supply either an instance type (for choosing CPU vs GPU) or an image URI."
    assert expected_msg in str(e)