create_optim_in_bwd_wrapper

torchtune.training.create_optim_in_bwd_wrapper(model: Module, optim_dict: Dict[Parameter, Optimizer]) → OptimizerInBackwardWrapper

创建在反向传播中运行优化器步骤的包装器。

参数:

• model (torch.nn.Module) – 包含正在优化的参数的模型。目前，假设所有正在优化的参数都属于单个顶级模型。将访问 model 的 named_parameters 属性，以查找正在优化的参数的参数名称。

• optim_dict (Dict[torch.nn.Parameter, torch.optim.Optimizer]) – 从参数到优化器的映射。

返回:

在反向传播中运行的优化器状态的包装器。

返回类型:

OptimizerInBackwardWrapper