(
*, out_block_type: str, num_groups_out: int, embed_dim: int, out_channels: int, act_fn: str, fc_dim: int
)
| 693 | |
| 694 | |
| 695 | def get_out_block( |
| 696 | *, out_block_type: str, num_groups_out: int, embed_dim: int, out_channels: int, act_fn: str, fc_dim: int |
| 697 | ) -> Optional[OutBlockType]: |
| 698 | if out_block_type == "OutConv1DBlock": |
| 699 | return OutConv1DBlock(num_groups_out, out_channels, embed_dim, act_fn) |
| 700 | elif out_block_type == "ValueFunction": |
| 701 | return OutValueFunctionBlock(fc_dim, embed_dim, act_fn) |
| 702 | return None |
no test coverage detected