diff --git a/hft.py b/hft.py
index ef4bfeb..31bdac9 100644
--- a/hft.py
+++ b/hft.py
@@ -9,7 +9,7 @@ import pytesseract
 import docx2txt
 import PyPDF2
 import json
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 from collections import defaultdict
 from huggingface_hub import login
 
@@ -21,8 +21,8 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 
 def free_memory():
-    torch.cuda.empty_cache()
-    torch.cuda.ipc_collect()
+    torch.empty_cache('cuda')
+    torch.ipc_collect('cuda')
 
 class SourceMapper:
     def __init__(self):