From 67f729ae68290761b40d54a8ecd620afd4649f9e Mon Sep 17 00:00:00 2001
From: "l.gabrysiak" <l.gabrysiak@osadkowski.pl>
Date: Tue, 25 Feb 2025 15:11:33 +0100
Subject: [PATCH] mod

---
 hft.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

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