autocast.encoders.unet# class UNetEncoder(in_channels, out_channels, cond_channels=0, hid_channels=(64, 128, 256), hid_blocks=(3, 3, 3), kernel_size=3, stride=2, attention_heads=None, spatial=2, periodic=False, identity_init=False, **kwargs)[source]# Bases: Encoder Base encoder. Parameters: in_channels (int) out_channels (int) cond_channels (int) hid_channels (Sequence[int]) hid_blocks (Sequence[int]) kernel_size (int | Sequence[int]) stride (int | Sequence[int]) attention_heads (dict[int, int] | None) spatial (int) periodic (bool) identity_init (bool) encoder_model: Module# encode(batch)[source]# Encode the input tensor into the latent space. Parameters: batch (Batch) – Input batch to be encoded. Returns: Encoded tensor in the latent space with shape (B, *, C_latent). Return type: Tensor
