tiny_diff.models.unet.layers.UNetLayerMixin

class tiny_diff.models.unet.layers.UNetLayerMixin

Bases: object

Mixin for UNet’s layers without skip connections.

__init__()

Methods

__init__()
forward(x, *args, **kwargs)	Forward pass.
residual_forward(h, *args, **kwargs)	Forward of the residual blocks.

Attributes

residual_block_cls

Class to use in residual blocks.