bagua.torch_api.bucket¶

Module Contents¶

class bagua.torch_api.bucket.BaguaBucket(tensors, name, flatten, alignment=1)¶

Create a Bagua bucket with a list of Bagua tensors.

Parameters

tensors – A list of Bagua tensors to be put in the bucket.
name – The unique name of the bucket.
flatten – If True, flatten the input tensors so that they are contiguous in memory.
alignment – If alignment > 1, Bagua will create a padding tensor to the bucket so that the total number of elements in the bucket divides the given alignment.

check_flatten(self)¶