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,
)