Function save_feed_forward

python/unet.py:50–57

(feed_forward, path)

    save_linear(geglu.proj, os.path.join(path, 'proj'))

def save_feed_forward(feed_forward, path):
    pathlib.Path(path).mkdir(parents=True, exist_ok=True)

    # Save GEGLU module
    save_geglu(feed_forward.net[0], os.path.join(path, 'geglu'))

    # Save Linear layer
    save_linear(feed_forward.net[2], os.path.join(path, 'lin'))


def save_basic_transformer_block(basic_transformer_block, path):