Convert primitive modules to the given dtype.
(l, dtype)
| 35 | |
| 36 | |
| 37 | def convert_module_to(l, dtype): |
| 38 | """ |
| 39 | Convert primitive modules to the given dtype. |
| 40 | """ |
| 41 | if isinstance(l, MIX_PRECISION_MODULES): |
| 42 | for p in l.parameters(): |
| 43 | p.data = p.data.to(dtype) |
| 44 | |
| 45 | |
| 46 | def zero_module(module): |