communication.serializer_plugins.torch module

(De)serialization logic for pytorch objects.

communication.serializer_plugins.torch.register()

Register pytorch serializer and deserializer.

Return type:

None

communication.serializer_plugins.torch.torch_deserialize_tensor(obj, **_kwargs)

Function for deserializing pytorch tensors and loading them into cpu

Parameters:

• obj (dict[str, Any] | bytes) – pytorch tensor to deserialize

• **_kwargs (Any) – optional extra keyword arguments

Return type:

Tensor

Returns:

deserialized tensor

communication.serializer_plugins.torch.torch_serialize_tensor(obj, **_kwargs)

Function for serializing pytorch tensors

Parameters:

• obj (Tensor) – pytorch object to serialize

• **_kwargs (Any) – optional extra keyword arguments

Return type:

tuple[int, ...] | bytes

Returns:

serialized tensor