diff --git a/hft.py b/hft.py
index 8909b4b..8a3486c 100644
--- a/hft.py
+++ b/hft.py
@@ -12,9 +12,6 @@ import json
 from collections import defaultdict
 from huggingface_hub import login
 
-import torch
-torch.cuda.empty_cache()
-
 login(token="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")
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
@@ -164,7 +161,7 @@ tokenized_dataset = dataset.map(tokenize_function, batched=True, batch_size=8)
 # Inicjalizacja modelu
 config = AutoModelForCausalLM.from_pretrained(model_name).config
 model = CustomModel.from_pretrained(model_name, config=config)
-model.gradient_checkpointing_enable()
+model.to("cpu")
 
 # Konfiguracja treningu
 training_args = TrainingArguments(
@@ -177,8 +174,7 @@ training_args = TrainingArguments(
     logging_steps=100,
     save_strategy="steps",
     save_steps=1000,
-    report_to="none",
-    gradient_checkpointing=True
+    report_to="none"
 )
 
 # Trening