PyTorch 加载模型出现 got an unexpected keyword argument ‘assign‘

PyTorch 加载模型出现 got an unexpected keyword argument ‘assign’

flyfish
错误提示

common\_base_model.py", line 433, in load
    model.load_state_dict(content["state_dict"], strict=True, assign=True)
TypeError: Module.load_state_dict() got an unexpected keyword argument 'assign'

查看PyTorch版本

import torch
print(torch.__version__) 
exit()

是2.0.0
查看源码
结果没有这个参数
https://github.com/pytorch/pytorch/blob/v2.0.0/torch/nn/modules/module.py

def load_state_dict(self, state_dict: Mapping[str, Any],
                    strict: bool = True):

看高版本的PyTorch版本有这个参数
https://github.com/pytorch/pytorch/blob/v2.3.0/torch/nn/modules/module.py

def load_state_dict(self, state_dict: Mapping[str, Any],
                    strict: bool = True, assign: bool = False):

如果使用assign这个参数需要升级下PyTorch版本