• Harry's avatar
    Remove all module wrapper's module when saving checkpoint (#399) · 5704613e
    Harry authored
    * fix: remove all module wrapper when saving checkpoint
    
    * refactor: move position of if
    
    * docs: add docstring
    
    * refactor: add _save_to_state_dict from official torch
    
    * docs: modify docstring of _save_to_state_dict
    
    * docs: modify docstring
    
    * feat: add unittest
    
    * feat: add DataParallel to unittest
    
    * fix: a bug when model has batchnorm
    
    * docs: update docstring
    5704613e
test_checkpoint.py 4.69 KB