ValueError: If you are using a dictionary as input, the data_keys should be None.

steps to reproduce this error:
```
1. log into CCC
2. python3 run_tests.py
```
error message:
```shell
_______ test_run_benchmark[configs/benchmark_v2_simple.yaml-False-False] _______

config = Namespace(defaults=Namespace(trainer_args={'precision': 'bf16-mixed', 'max_epochs': 10}, terratorch_task={'model_args'...lticlass_Jaccard_Index', early_prune=False, early_stop_patience=10, optimization_except=set(), max_run_duration=None)])
continue_existing_experiment = False, test_models = False

@pytest.mark.parametrize(
"config, continue_existing_experiment, test_models",
[
("configs/benchmark_v2_simple.yaml", False, False),
# ("configs/benchmark_v2_simple.yaml", False, True),
# ("configs/benchmark_v2_simple.yaml", True, True),
# ("configs/benchmark_v2_simple.yaml", True, False),
(
"configs/dofa_large_patch16_224_upernetdecoder_true_modified.yaml",
True,
True,
),
(
"configs/dofa_large_patch16_224_upernetdecoder_true_modified.yaml",
True,
False,
),
(
"configs/dofa_large_patch16_224_upernetdecoder_true_modified.yaml",
False,
True,
),
(
"configs/dofa_large_patch16_224_upernetdecoder_true_modified.yaml",
False,
False,
),
],
)
def test_run_benchmark(
config: str, continue_existing_experiment: bool, test_models: bool
):

path = os.path.join(os.getcwd(), config)
config_path = Path(path)
assert (
config_path.exists()
), f"Error! config does not exist: {config_path.resolve()}"
# instantiate objects from yaml
parser = ArgumentParser()
parser.add_argument('--defaults', type=Defaults) # to ignore model
parser.add_argument('--optimization_space', type=dict) # to ignore model
parser.add_argument('--experiment_name', type=str) # to ignore model
parser.add_argument('--run_name', type=str) # to ignore model
parser.add_argument('--save_models', type=bool) # to ignore model
parser.add_argument('--storage_uri', type=str) # to ignore model
parser.add_argument('--ray_storage_path', type=str) # to ignore model
parser.add_argument('--n_trials', type=int) # to ignore model
parser.add_argument('--run_repetitions', type=int) # to ignore model
parser.add_argument('--tasks', type=list[Task])
config = parser.parse_path(str(config_path))
config_init = parser.instantiate_classes(config)
# validate the objects
experiment_name = config_init.experiment_name
experiment_name = f"{experiment_name}_continue_{continue_existing_experiment}_test_models_{test_models}"
assert isinstance(experiment_name, str), f"Error! {experiment_name=} is not a str"
run_name = config_init.run_name
if run_name is not None:
assert isinstance(run_name, str), f"Error! {run_name=} is not a str"
tasks = config_init.tasks
assert isinstance(tasks, list), f"Error! {tasks=} is not a list"
for t in tasks:
assert isinstance(t, Task), f"Error! {t=} is not a Task"
defaults = config_init.defaults
assert isinstance(defaults, Defaults), f"Error! {defaults=} is not a Defaults"
# defaults.trainer_args["max_epochs"] = 5
storage_uri = OUTPUT_DIR
assert isinstance(storage_uri, str), f"Error! {storage_uri=} is not a str"
storage_uri_path = Path(storage_uri) / uuid.uuid4().hex / "hpo"
if not storage_uri_path.exists():
try:
storage_uri_path.mkdir(parents=True, exist_ok=True)
print(f"Directory created at: {path}")
except FileNotFoundError as e:
print(f"Error creating directory: {e}")
optimization_space = config_init.optimization_space
assert isinstance(
optimization_space, dict
), f"Error! {optimization_space=} is not a dict"
ray_storage = RAY_STORAGE
assert isinstance(ray_storage, str), f"Error! {ray_storage=} is not a str"
ray_storage_path = Path(ray_storage) / uuid.uuid4().hex
if not ray_storage_path.exists():
try:
ray_storage_path.mkdir(parents=True, exist_ok=True)
print(f"Directory created at: {path}")
except FileNotFoundError as e:
print(f"Error creating directory: {e}")
n_trials = config_init.n_trials
assert isinstance(n_trials, int) and n_trials > 0, f"Error! {n_trials=} is invalid"
# run_repetions is an optional parameter
run_repetitions = config_init.run_repetitions
if run_repetitions is not None:
assert (
isinstance(run_repetitions, int) and run_repetitions >= 0
), f"Error! {run_repetitions=} is invalid"
else:
run_repetitions = 0
> mlflow_experiment_id = benchmark_backbone(
experiment_name=experiment_name,
run_name=run_name,
run_id=None,
defaults=defaults,
tasks=tasks,
n_trials=n_trials,
save_models=False,
storage_uri=str(storage_uri_path),
ray_storage_path=str(ray_storage_path),
optimization_space=optimization_space,
continue_existing_experiment=continue_existing_experiment,
test_models=test_models,
run_repetitions=run_repetitions,
)

tests/test_benchmark.py:203:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
benchmark/backbone_benchmark.py:304: in benchmark_backbone
best_value, metric_name, hparams = benchmark_backbone_on_task(
benchmark/backbone_benchmark.py:127: in benchmark_backbone_on_task
study.optimize(
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/optuna/study/study.py:475: in optimize
_optimize(
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/optuna/study/_optimize.py:63: in _optimize
_optimize_sequential(
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/optuna/study/_optimize.py:160: in _optimize_sequential
frozen_trial = _run_trial(study, func, catch)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/optuna/study/_optimize.py:248: in _run_trial
raise func_err
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/optuna/study/_optimize.py:197: in _run_trial
value_or_values = func(trial)
benchmark/model_fitting.py:444: in fit_model_with_hparams
return fit_model(
benchmark/model_fitting.py:396: in fit_model
launch_training(
benchmark/model_fitting.py:272: in launch_training
trainer.fit(task, datamodule=datamodule)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py:561: in fit
call._call_and_handle_interrupt(
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/trainer/call.py:48: in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py:599: in _fit_impl
self._run(model, ckpt_path=ckpt_path)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py:1012: in _run
results = self._run_stage()
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py:1054: in _run_stage
self._run_sanity_check()
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/trainer/trainer.py:1083: in _run_sanity_check
val_loop.run()
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/loops/utilities.py:179: in _decorator
return loop_run(self, *args, **kwargs)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/loops/evaluation_loop.py:145: in run
self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/loops/evaluation_loop.py:411: in _evaluation_step
batch = call._call_strategy_hook(trainer, "batch_to_device", batch, dataloader_idx=dataloader_idx)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/trainer/call.py:328: in _call_strategy_hook
output = fn(*args, **kwargs)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/strategies/strategy.py:278: in batch_to_device
return model._apply_batch_transfer_handler(batch, device=device, dataloader_idx=dataloader_idx)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/core/module.py:353: in _apply_batch_transfer_handler
batch = self._call_batch_hook("on_after_batch_transfer", batch, dataloader_idx)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/core/module.py:341: in _call_batch_hook
return trainer_method(trainer, hook_name, *args)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/lightning/pytorch/trainer/call.py:198: in _call_lightning_datamodule_hook
return fn(*args, **kwargs)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/torchgeo/datamodules/geo.py:141: in on_after_batch_transfer
batch = aug(batch)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/kornia/augmentation/container/augment.py:516: in __call__
_output_image = decorated_forward(*inputs, **kwargs)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/kornia/core/module.py:81: in wrapper
tensor_outputs = func(*args, **kwargs)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/torch/nn/modules/module.py:1739: in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/torch/nn/modules/module.py:1750: in _call_impl
return forward_call(*args, **kwargs)
/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/kornia/augmentation/container/augment.py:439: in forward
original_keys, data_keys, args, invalid_data = self._preproc_dict_data(args[0])
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = AugmentationSequential(
(Normalize_0): Normalize(p=1.0, p_batch=1.0, same_on_batch=True, mean=tensor([0.4570, 0.5201, 0.4808, 0.5699]), std=tensor([0.2119, 0.1977, 0.1744, 0.2832]))
)
data = {'image': tensor([[[[0.0000, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.0000, ..., 0... [2, 2, 2, ..., 2, 2, 2],
[2, 2, 2, ..., 2, 2, 2],
[2, 2, 2, ..., 2, 2, 2]]], device='cuda:0')}

def _preproc_dict_data(
self, data: Dict[str, DataType]
) -> Tuple[Tuple[str, ...], List[DataKey], Tuple[DataType, ...], Optional[Dict[str, Any]]]:
if self.data_keys is not None:
> raise ValueError("If you are using a dictionary as input, the data_keys should be None.")
E ValueError: If you are using a dictionary as input, the data_keys should be None.

/u/ltizzei/.pyenv/versions/iterate/lib/python3.12/site-packages/kornia/augmentation/container/augment.py:547: ValueError
----------------------------- Captured stdout call -----------------------------
Directory created at: /u/ltizzei/Projects/Orgs/IBM/terratorch-iterate/configs/benchmark_v2_simple.yaml
Directory created at: /u/ltizzei/Projects/Orgs/IBM/terratorch-iterate/configs/benchmark_v2_simple.yaml

Sanity Checking: | | 0/? [00:00