Running --mode test Returns Error

[kvarys]

Dear Researchers,

When executing python3 main.py --mode test --cvae_path craft_cvae_abstract.pth --config craft_ours.yaml --test_iter 400000 --viz_name viz I receive the following error:

'Trainer' object has no attribute 'scenario_pool' -> trainers/partial.py(791)test() -> possible_tasks = list(self.scenario_pool.keys())

At line 812 in the same method I noticed the following call which initializes the scenario_pool:

if self.world_name == "CraftWorldHard":                                                           
     self._load_or_sample_test_scenarios(50)                                                                     
else:                                                                                                       
     self._load_or_sample_test_scenarios(10)

by adding this if/else into the test method before the line 791, the problem seems to be resolved. Would you say this is a correct solution of the issue?

Thank you!