def test_bad_stop_training_val(): nn_state = PositiveWaveFunction(10, gpu=False) msg = "Setting stop_training to a non-boolean value should have raised an error." with pytest.raises(ValueError): nn_state.stop_training = "foobar" pytest.fail(msg)
def test_stop_training(gpu): qucumber.set_random_seed(SEED, cpu=True, gpu=gpu, quiet=True) nn_state = PositiveWaveFunction(10, gpu=gpu) old_params = parameters_to_vector(nn_state.rbm_am.parameters()) data = torch.ones(100, 10) nn_state.stop_training = True nn_state.fit(data) new_params = parameters_to_vector(nn_state.rbm_am.parameters()) msg = "stop_training didn't work!" assert torch.equal(old_params, new_params), msg