Converts a module's immediate parameters and buffers to dtype.
(module, dtype)
| 59 | |
| 60 | |
| 61 | def convert_module(module, dtype): |
| 62 | """ |
| 63 | Converts a module's immediate parameters and buffers to dtype. |
| 64 | """ |
| 65 | for param in module.parameters(recurse=False): |
| 66 | if param is not None: |
| 67 | if param.data.dtype.is_floating_point: |
| 68 | param.data = param.data.to(dtype=dtype) |
| 69 | if param._grad is not None and param._grad.data.dtype.is_floating_point: |
| 70 | param._grad.data = param._grad.data.to(dtype=dtype) |
| 71 | |
| 72 | for buf in module.buffers(recurse=False): |
| 73 | if buf is not None and buf.data.dtype.is_floating_point: |
| 74 | buf.data = buf.data.to(dtype=dtype) |
| 75 | |
| 76 | |
| 77 | def convert_network(network, dtype): |