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