(cast_fn, x, cache)
| 133 | return new_args |
| 134 | |
| 135 | def cached_cast(cast_fn, x, cache): |
| 136 | if is_nested(x): |
| 137 | return type(x)([cached_cast(y) for y in x]) |
| 138 | if x in cache: |
| 139 | cached_x = cache[x] |
| 140 | # During eval, it's possible to end up caching casted weights |
| 141 | # with requires_grad == False. This is then a problem when they |
| 142 | # get reused on the next train iter. So we ensure that cached |
| 143 | # weights have same requires_grad flag of most recent request. |
| 144 | if x.requires_grad != cached_x.requires_grad: |
| 145 | cached_x.requires_grad_(x.requires_grad) |
| 146 | return cache[x] |
| 147 | |
| 148 | casted_x = cast_fn(x) |
| 149 | cache[x] = casted_x |
| 150 | return casted_x |
| 151 | |
| 152 | def verbosify(cast_fn, fn_name, verbose): |
| 153 | if verbose: |
nothing calls this directly
no test coverage detected