This commit is contained in:
l.gabrysiak 2025-02-25 15:11:33 +01:00
parent 4f8d57e770
commit 67f729ae68
1 changed files with 3 additions and 3 deletions

6
hft.py
View File

@ -9,7 +9,7 @@ import pytesseract
import docx2txt import docx2txt
import PyPDF2 import PyPDF2
import json import json
from torch.cuda.amp import autocast from torch.amp import autocast
from collections import defaultdict from collections import defaultdict
from huggingface_hub import login from huggingface_hub import login
@ -21,8 +21,8 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
def free_memory(): def free_memory():
torch.cuda.empty_cache() torch.empty_cache('cuda')
torch.cuda.ipc_collect() torch.ipc_collect('cuda')
class SourceMapper: class SourceMapper:
def __init__(self): def __init__(self):