(cls, origin_dtype, *inputs, **kwargs)
| 34 | param.data = param.data.to(torch.float8_e4m3fn) |
| 35 | |
| 36 | def autocast_model_forward(cls, origin_dtype, *inputs, **kwargs): |
| 37 | weight_dtype = cls.weight.dtype |
| 38 | cls.to(origin_dtype) |
| 39 | |
| 40 | # Convert all inputs to the original dtype |
| 41 | inputs = [input.to(origin_dtype) for input in inputs] |
| 42 | out = cls.original_forward(*inputs, **kwargs) |
| 43 | |
| 44 | cls.to(weight_dtype) |
| 45 | return out |
| 46 | |
| 47 | def convert_weight_dtype_wrapper(module, origin_dtype): |
| 48 | for name, module in module.named_modules(): |
no outgoing calls
no test coverage detected