(spatial_transformer, path)
| 72 | |
| 73 | |
| 74 | def save_spatial_transformer(spatial_transformer, path): |
| 75 | pathlib.Path(path).mkdir(parents=True, exist_ok=True) |
| 76 | |
| 77 | # Save GroupNorm, Conv2d, BasicTransformerBlock instances |
| 78 | save_group_norm(spatial_transformer.norm, os.path.join(path, 'norm')) |
| 79 | save_conv2d(spatial_transformer.proj_in, os.path.join(path, 'proj_in')) |
| 80 | save_basic_transformer_block(spatial_transformer.transformer_blocks[0], os.path.join(path, 'transformer')) |
| 81 | save_conv2d(spatial_transformer.proj_out, os.path.join(path, 'proj_out')) |
| 82 | |
| 83 | |
| 84 | def save_downsample(downsample, path): |
no test coverage detected