Function AutoEncoder

core/models/autoencoder.py:86–93

(img_size)

def AutoEncoder(img_size):
    assert isinstance(img_size, tuple), f"img_size should be tuple, but got {type(img_size)}"
    if img_size == (1, 28, 28):
        return AutoEncoder1x28x28()
    elif img_size == (3, 32, 32):
        return AutoEncoder3x32x32()
    else:
        raise NotImplementedError("Unsupported img size!")


if __name__=='__main__':