Many thanks for your fantastic project. Just a kind suggestion: while using bmtrain, it will throw out a lot warnings and they are all corresponding to torch.storage as following:

/opt/conda/envs/compression/lib/python3.10/site-packages/bmtrain/store.py:178: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly. To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()

it seems that it is related to feature of torch. You may replace all the tensor.storage() with tensor.untyped_storage() so that we can get the clean logs:)