train_on_ray2 Tuner Error

train_on_ray2.ipynb

Can you try removing datasets from param_space?

param_space = {
    "train_loop_config" : {
        "lr": tune.grid_search([0.001, 0.01]),
        "batch_size": batch_size,
        "image_shape": image_shape,
        "num_epochs": 10,
    },
    "datasets": {"train": train_ds, "val": val_ds}
}

basically the above blob becomes

param_space = {
    "train_loop_config" : {
        "lr": tune.grid_search([0.001, 0.01]),
        "batch_size": batch_size,
        "image_shape": image_shape,
        "num_epochs": 10,
    },
}

The idea is that if you don't need to tune the dataset, you only need to supply it through Trainer() args as you did. No need to do it through param_space again. Actually similarly for all the none-tuning parameters.