def bufferize(worker: AbstractWorker, script_module: torch.jit.ScriptModule) -> ScriptModulePB:
"""
This method serializes a torch.jit.ScriptModule using ScriptModulePB.
Args:
script_module (torch.jit.ScriptModule): input jit.ScriptModule to be serialized.
Returns:
protobuf_script (ScriptModulePB): serialized jit.ScriptModule.
"""
protobuf_script = ScriptModulePB()
protobuf_script.obj = script_module.save_to_buffer()
return protobuf_script
def _simplify_script_module(obj: torch.jit.ScriptModule) -> str:
"""Strategy to serialize a script module using Torch.jit"""
return obj.save_to_buffer()
def _simplify_script_module(worker: AbstractWorker, obj: torch.jit.ScriptModule) -> Tuple:
"""Strategy to serialize a script module using Torch.jit"""
return (obj.save_to_buffer(),)
def _bufferize_script_module(
worker: AbstractWorker,
script_module: torch.jit.ScriptModule) -> ScriptModulePB:
protobuf_script = ScriptModulePB()
protobuf_script.obj = script_module.save_to_buffer()
return protobuf_script