diff --git a/README.md b/README.md
index a937856..3c2b54f 100644
--- a/README.md
+++ b/README.md
@@ -15,7 +15,7 @@ For more details, please refer to our paper:
 
 *Liang Jin\*, [Jiancheng Yang](http://jiancheng-yang.com/)\*, [Kaiming Kuang](http://kaimingkuang.github.io/), [Bingbing Ni](https://scholar.google.com/citations?user=eUbmKwYAAAAJ), Yiyi Gao, Yingli Sun, Pan Gao, Weiling Ma, Mingyu Tan, Hui Kang, Jiajun Chen, Ming Li*
 
-EBioMedicine ([DOI](https://doi.org/10.1016/j.ebiom.2020.103106))
+EBioMedicine, 2020 ([DOI](https://doi.org/10.1016/j.ebiom.2020.103106))
 
 ## Code Structure
 * FracNet/
@@ -76,6 +76,10 @@ python -m predict --image_dir <image_directory> --pred_dir <predition_directory>
 
 In the [predict.py](predict.py), we adopt a post-processing procedure of [removing low-probability regions](predict.py#L18), [spine regions](predict.py#L24), and [small objects](predict.py#L48). This procedure leads to fewer false negatives. You may also skip the post-processing by setting `--postprocess False` in the command line argument and check the raw output.
 
+***Note 1***: This project aims at a prototype for RibFrac Challenge; However, as the challenge data provider, we would like to avoid unintended data leakage. Therefore, we did **NOT** provide all details for the models, including those in both training and inference stages. Nevertheless, it is guaranteed that the performance in the EBioMedicine'20 paper could be reproduced with this one-stage FracNet using 3D UNet as backbone, but also note that 1) it is trained with challenge data and extra in-house data, 2) there is heavy pre-processing and post-processing (including *rib segmentation*). See more discussion in this [issue](https://github.com/M3DV/FracNet/issues/7).
+
+***Note 2***: Our paper on rib segmentation and centerline extraction has been accepted by [MICCAI'21](https://arxiv.org/abs/2109.09521). Feel free to check it ;)
+
 ### Evaluation
 To evaluate your prediction, run the following in command line:
 ```bash
diff --git a/main.py b/main.py
index 3ed479b..42f4105 100644
--- a/main.py
+++ b/main.py
@@ -1,6 +1,7 @@
 from functools import partial
 
-import torch.nn as nn
+from torch import save
+
 from fastai.basic_train import Learner
 from fastai.train import ShowGraph
 from fastai.data_block import DataBunch
@@ -67,7 +68,7 @@ def main(args):
     )
 
     if args.save_model:
-        torch.save(model.module.state_dict(), "./model_weights.pth")
+        save(model.module.state_dict(), "./model_weights.pth")
 
 
 if __name__ == "__main__":