mod

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):