Source code for syft.lib.torch.size

# third party
import torch

# syft relative
from ...generate_wrapper import GenerateWrapper
from ...proto.lib.torch.size_pb2 import Size as TorchSize_PB


[docs]def protobuf_torch_size_serializer(torch_size: torch.Size) -> TorchSize_PB:
    serialized_size = TorchSize_PB(data=torch_size)
    return serialized_size


[docs]def protobuf_torch_size_deserializer(proto_size: TorchSize_PB) -> torch.Size:
    torch_size = torch.Size(proto_size.data)
    return torch_size


GenerateWrapper(
    wrapped_type=torch.Size,
    import_path="torch.Size",
    protobuf_scheme=TorchSize_PB,
    type_object2proto=protobuf_torch_size_serializer,
    type_proto2object=protobuf_torch_size_deserializer,
)