Converts a network's parameters and buffers to dtype.
(network, dtype)
| 75 | |
| 76 | |
| 77 | def convert_network(network, dtype): |
| 78 | """ |
| 79 | Converts a network's parameters and buffers to dtype. |
| 80 | """ |
| 81 | for module in network.modules(): |
| 82 | if isinstance(module, torch.nn.modules.batchnorm._BatchNorm) and module.affine is True: |
| 83 | continue |
| 84 | convert_module(module, dtype) |
| 85 | return network |
| 86 | |
| 87 | |
| 88 | class FP16Model(nn.Module): |
no test coverage detected