会将Tensor拷贝到GPU上,避免模型训练过程中频繁的H2D操作,默认persistent=True 默认torch.save时将register_buffer中的tensor也写入dict, 可能导致模型load.state.dict(strict = False)时出现key错误。如果persistent=False则不会写书state_dict
pytorch register_buffer(name, tensor, persistent=True)
最新推荐文章于 2024-07-25 00:38:11 发布