def METHOD_NAME(self):
    @find_executable_batch_size(starting_batch_size=16)
    def mock_training_loop_function(batch_size):
        raise ValueError("Oops, we had an error!")
    with self.assertRaises(ValueError) as cm:
        mock_training_loop_function()
        self.assertIn("Oops, we had an error!", cm.exception.args[0])