| 40 | return wrapper |
| 41 | |
| 42 | def cast_output(output, input_type, verbose=False): |
| 43 | if isinstance(output, dict): |
| 44 | keys = output.keys() |
| 45 | for k in keys: |
| 46 | output[k] = cast_output(output[k], input_type) |
| 47 | return output |
| 48 | |
| 49 | if utils.is_fp_tensor(output) and output.type() != input_type: |
| 50 | if verbose: |
| 51 | print('ori output type: {}, input type: {}'.format(output.type(), input_type)) |
| 52 | return output.type(input_type) |
| 53 | return output |
| 54 | |
| 55 | def cached_cast(mod, fn, cast_fn, handle, |
| 56 | try_caching=False, verbose=False): |