From 0ace5d134830b578a9fb54de084128be442902c1 Mon Sep 17 00:00:00 2001 From: "l.gabrysiak" Date: Tue, 25 Feb 2025 14:50:09 +0100 Subject: [PATCH] mod --- hft.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/hft.py b/hft.py index 756535e..d1e432a 100644 --- a/hft.py +++ b/hft.py @@ -10,6 +10,7 @@ import pytesseract import docx2txt import PyPDF2 import json +from torch.cuda.amp import autocast from collections import defaultdict from huggingface_hub import login @@ -165,6 +166,8 @@ tokenized_dataset = dataset.map(tokenize_function, batched=True, batch_size=32) config = AutoModelForCausalLM.from_pretrained(model_name).config #model = CustomModel.from_pretrained(model_name, config=config) model = CustomModel.from_pretrained(model_name) +model.config.gradient_checkpointing = True +model.config.use_cache = False model.resize_token_embeddings(len(tokenizer)) model.gradient_checkpointing_enable()