diff --git a/examples/neptune_plugin/neptune_plugin/neptune_example.py b/examples/neptune_plugin/neptune_plugin/neptune_example.py index f74d2e580..b6a7949b8 100644 --- a/examples/neptune_plugin/neptune_plugin/neptune_example.py +++ b/examples/neptune_plugin/neptune_plugin/neptune_example.py @@ -62,6 +62,7 @@ # %% [markdown] # First, we use a task to download the dataset and cache the data in Flyte: + # %% @task( container_image=image, @@ -85,6 +86,7 @@ def get_dataset() -> Tuple[np.ndarray, np.ndarray]: # in Neptune's integrations. In this example, we pass the `Run` object into Neptune's # XGBoost callback. + # %% @task( container_image=image, @@ -133,6 +135,7 @@ def train_model(max_depth: int, X: np.ndarray, y: np.ndarray): # With Flyte's dynamic workflows, we can scale up multiple training jobs with different # `max_depths`: + # %% @dynamic(container_image=image) def train_multiple_models(max_depths: List[int], X: np.ndarray, y: np.ndarray):