(net_, count, place_in_unet)
| 220 | controller = DummyController() |
| 221 | |
| 222 | def register_recr(net_, count, place_in_unet): |
| 223 | if net_.__class__.__name__ == 'CrossAttention': |
| 224 | net_.forward = ca_forward(net_, place_in_unet) |
| 225 | return count + 1 |
| 226 | elif hasattr(net_, 'children'): |
| 227 | for net__ in net_.children(): |
| 228 | count = register_recr(net__, count, place_in_unet) |
| 229 | return count |
| 230 | |
| 231 | cross_att_count = 0 |
| 232 | sub_nets = model.unet.named_children() |
no test coverage detected