def METHOD_NAME(self):
    torch.manual_seed(0)
    model = UNet2DModel(
        sample_size=(32, 64),
        in_channels=1,
        out_channels=1,
        layers_per_block=2,
        block_out_channels=(128, 128),
        down_block_types=("AttnDownBlock2D", "DownBlock2D"),
        up_block_types=("UpBlock2D", "AttnUpBlock2D"),
    )
    return model