Function save_resnet_block

python/autoencoder.py:7–16

(resnet_block, path)

from tinygrad.nn import Conv2d

def save_resnet_block(resnet_block, path):
    pathlib.Path(path).mkdir(parents=True, exist_ok=True)
    
    save_group_norm(resnet_block.norm1, pathlib.Path(path, 'norm1'))
    save_conv2d(resnet_block.conv1, pathlib.Path(path, 'conv1'))
    save_group_norm(resnet_block.norm2, pathlib.Path(path, 'norm2'))
    save_conv2d(resnet_block.conv2, pathlib.Path(path, 'conv2'))

    if isinstance(resnet_block.nin_shortcut, Conv2d):
        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)