Fix for offloading when using TorchAO >= 0.7.0 (#3332)
* fix
* update
* fix
* apply suggestions from review
Co-Authored-By: Benjamin Bossan <BenjaminBossan@users.noreply.github.com>
Co-Authored-By: Xuehai Pan <XuehaiPan@pku.edu.cn>
* make style
---------
Co-authored-by: Xuehai Pan <XuehaiPan@pku.edu.cn>
