[FSDP] Address follow up comments for CPU offload (#67813)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/67813
Address Shen's comments in
https://github.com/pytorch/pytorch/pull/67249/files
ghstack-source-id: 142379312
Test Plan: CI
Reviewed By: zhaojuanmao
Differential Revision: D32157545
fbshipit-source-id: 3cc2df6d5fa0d3b9383ed3711e7f79729dbb1dda
