Get the class of a block from a model, using the block's class name.
(model: torch.nn.Module, block_class_name: str)
| 118 | |
| 119 | |
| 120 | def get_block_class_from_model(model: torch.nn.Module, block_class_name: str) -> torch.nn.Module: |
| 121 | """Get the class of a block from a model, using the block's class name.""" |
| 122 | for module in model.modules(): |
| 123 | if module.__class__.__name__ == block_class_name: |
| 124 | return module.__class__ |
| 125 | raise ValueError(f"Could not find block class {block_class_name} in model {model}") |
| 126 | |
| 127 | |
| 128 | def get_block_class_from_model_class_and_block_name(model_class: Type, block_class_name: str) -> Type: |