Only make a shallow copy when loading optimizer state_dict (#106082)

The thing we do still deep copy is the param_groups, which is much lighter weight. This should also save memory when loading from a checkpoint.

The deepcopy was introduced in ecfcf39f30, but module.py had only a shallow copy at that point so it did not actually bring parity.

Incorporates an XLA fix, which is why I'm updating the pin to ca5eab87a7

Pull Request resolved: https://github.com/pytorch/pytorch/pull/106082
Approved by: https://github.com/albanD, https://github.com/Skylion007
This commit is contained in:
Jane Xu
2023-07-31 19:23:08 -07:00
committed by PyTorch MergeBot
parent ceea08a986
commit 59d0dea90f
4 changed files with 10 additions and 8 deletions

View File

@ -36,3 +36,5 @@ dd3a77bc965adf9fe8ba582ee13bb7f14c9661b0
f70844bec783bfce43c950ccf180dc494e86f2bf
# 2023-07-28 Apply UFMT to all non test/torch files
e6ec0efaf87703c5f889cfc20b29be455885d58d
# 2023-07-31 [optim][BE] split test file into logical parts: SWA, LR, optim
a53cda1ddc15336dc1ff0ce1eff2a49cdc5f882e