Function save_attn_block

python/autoencoder.py:18–25

(attn_block, path)

        save_conv2d(resnet_block.nin_shortcut, pathlib.Path(path, 'nin_shortcut'))

def save_attn_block(attn_block, path):
    pathlib.Path(path).mkdir(parents=True, exist_ok=True)
    
    save_group_norm(attn_block.norm, pathlib.Path(path, 'norm'))
    save_conv2d(attn_block.q, pathlib.Path(path, 'q'))
    save_conv2d(attn_block.k, pathlib.Path(path, 'k'))
    save_conv2d(attn_block.v, pathlib.Path(path, 'v'))
    save_conv2d(attn_block.proj_out, pathlib.Path(path, 'proj_out'))

def save_mid(mid, path):
    pathlib.Path(path).mkdir(parents=True, exist_ok=True)