protectai / modelscan

Protection against Model Serialization Attacks

Home Page:http://modelscan.ai

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Pytorch tests

swashko opened this issue · comments

Add tests for pytorch new and old formats

New format tests (safe and unsafe files) added in #87
Old format unsafe file test still needed

Barrier to old format unsafe file test is issues with torch.save()
torch.save( torch.load(model_path), f=unsafe_file_path, pickle_module=PickleInject([payload]), _use_new_zipfile_serialization=False, )
gives an error of PickleInject needing a dump method to call
pickle_module.dump(MAGIC_NUMBER, f, protocol=pickle_protocol)
but adding a dump method results in an error that it received multiple values for protocol