• drhead's avatar
    Add extra performance features for EMAModel, torch._foreach operations and... · 2ada094b
    drhead authored
    
    Add extra performance features for EMAModel, torch._foreach operations and better support for non-blocking CPU offloading (#7685)
    
    * Add support for _foreach operations and non-blocking to EMAModel
    
    * default foreach to false
    
    * add non-blocking EMA offloading to SD1.5 T2I example script
    
    * fix whitespace
    
    * move foreach to cli argument
    
    * linting
    
    * Update README.md re: EMA weight training
    
    * correct args.foreach_ema
    
    * add tests for foreach ema
    
    * code quality
    
    * add foreach to from_pretrained
    
    * default foreach false
    
    * fix linting
    
    ---------
    Co-authored-by: default avatarSayak Paul <spsayakpaul@gmail.com>
    Co-authored-by: default avatardrhead <a@a.a>
    2ada094b
training_utils.py 24.5 KB