Return best model in Pytorch ?

[TheJproject]

Hello everyone,
I am trying NNI for the first time and I would like to return the best model from a NAS experiment. I have tried to get the model using
for model_dict in exp.export_top_models(formatter='dict'):
print(model_dict)
with model_context(model_dict):
final_model = Net()
But it seems that the model_context is only available in v3.0 ? I saw on old example that it is possible to return the Pytorch code.
I have also tried to export the nas_model output using nnictl command but I cannot figure out how to turn the dict output into a Pytorch model.
I am a bit out of ideas if anyone have any suggestions that would greatly appreciated. Thanks !