def bufferize(worker: AbstractWorker, script_module: torch.jit.ScriptFunction) -> ScriptFunctionPB: """ This method serializes a torch.jit.ScriptFunction into a ScriptFunctionPB. Args: script_module (torch.jit.ScriptFunction): input torch.jit.ScriptFunction to be serialized. Returns: protobuf_script (ScriptFunctionPB): serialized torch.jit.ScriptFunction. """ protobuf_script = ScriptFunctionPB() protobuf_script.obj = script_module.save_to_buffer() return protobuf_script
def _bufferize_script_function( worker: AbstractWorker, script_module: torch.jit.ScriptFunction) -> ScriptFunctionPB: protobuf_script = ScriptFunctionPB() protobuf_script.obj = script_module.save_to_buffer() return protobuf_script