Export model issue for netAug

[DasMadhureeta]

Hi,

While exporting the tiny model from a large model, I found there was an accuracy drop. Can you clarify this part? Is there anything that needs to do before exporting the model once the training is done? For this part- if use_netaug: checkpoint = load_state_dict_from_file(os.path.join(checkpoint_path, "best.pt")) model.module.load_state_dict(checkpoint) model.eval() model.module.set_active(mode="min") with torch.no_grad(): model.module( torch.zeros( 1, 3, exp_config["data_provider"]["image_size"], exp_config["data_provider"]["image_size"], ).cuda() ) export_model = model.module.export() if dist.is_master(): torch.save( {"state_dict": export_model.state_dict()}, os.path.join(checkpoint_path, "target.pt"), _use_new_zipfile_serialization=False, )