Function save_basic_transformer_block

python/unet.py:60–70

(basic_transformer_block, path)

def save_basic_transformer_block(basic_transformer_block, path):
    pathlib.Path(path).mkdir(parents=True, exist_ok=True)

    # Save CrossAttention, FeedForward and LayerNorm instances
    save_cross_attention(basic_transformer_block.attn1, os.path.join(path, 'attn1'))
    save_feed_forward(basic_transformer_block.ff, os.path.join(path, 'mlp'))
    save_cross_attention(basic_transformer_block.attn2, os.path.join(path, 'attn2'))

    save_layer_norm(basic_transformer_block.norm1, os.path.join(path, 'norm1'))
    save_layer_norm(basic_transformer_block.norm2, os.path.join(path, 'norm2'))  
    save_layer_norm(basic_transformer_block.norm3, os.path.join(path, 'norm3'))