modelzoo.vision.pytorch.dit.layers.vae.UNetMidBlock2D.UNetMidBlock2D

class modelzoo.vision.pytorch.dit.layers.vae.UNetMidBlock2D.UNetMidBlock2D

Bases: torch.nn.Module

Methods

forward

__call__(*args: Any, **kwargs: Any) → Any

Call self as a function.

__init__(in_channels: int, temb_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-06, resnet_time_scale_shift: str = 'default', resnet_act_fn: str = 'swish', resnet_groups: int = 32, resnet_pre_norm: bool = True, add_attention: bool = True, attn_num_head_channels=1, output_scale_factor=1.0, attention_type='aiayn_attention', extra_attn_params=None)

static __new__(cls, *args: Any, **kwargs: Any) → Any

previous

modelzoo.vision.pytorch.dit.layers.vae.UNetMidBlock2D

next

modelzoo.vision.pytorch.dit.layers.vae.UpDecoderBlock2D