Issue with AdamW on Apple M1

[sigma-andex]

There seems to be a problem when trying to run the denoising diffusion models on Apple Silicon. This seems to be the related keras issue: keras-team/tf-keras#176
See also: https://developer.apple.com/forums/thread/729732

I was able to get it running using the legacy.Adam optimizer:

ddm.compile(
    # optimizer=optimizers.experimental.AdamW(
    #     learning_rate=LEARNING_RATE, weight_decay=WEIGHT_DECAY
    # ),
    optimizer=optimizers.legacy.Adam(
        learning_rate=LEARNING_RATE
    ),
    loss=losses.mean_absolute_error,
)

Not sure how this effects the quality of this model though...