modelzoo.vision.pytorch.dit.layers.vae.utils.randn_tensor#

modelzoo.vision.pytorch.dit.layers.vae.utils.randn_tensor(shape: Union[Tuple, List], generator: Optional[Union[List[torch.Generator], torch.Generator]] = None, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, layout: Optional[torch.layout] = None)[source]#

This is a helper function that allows to create random tensors on the desired device with the desired dtype. When passing a list of generators one can seed each batched size individually. If CPU generators are passed the tensor will always be created on CPU.