Update README.md

README.md

@@ -37,6 +37,7 @@ import torch
 from transformers.modeling_outputs import BaseModelOutput
 from transformers import T5TokenizerFast
 import torch.nn.functional as F
+import shutil
 
 model_name = 'vamsibanda/sbert-onnx-all-roberta-large-v1'
 cache_folder = './'
@@ -49,7 +50,7 @@ def mean_pooling(model_output, attention_mask):
 
 def download_onnx_model(model_name, cache_folder, model_path, force_download = False):
    if force_download and os.path.exists(model_path):
-      os.remove(model_path)
+      shutil.rmtree(model_path)
   
    snapshot_download(model_name,
                      cache_dir=cache_folder,