From 5786267144851b695ac6d7c4fb41814cc9694818 Mon Sep 17 00:00:00 2001
From: randaller <randaller@users.noreply.github.com>
Date: Sun, 19 Mar 2023 13:25:03 +0300
Subject: [PATCH] Update hf-training-example.py

---
 hf-training-example.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/hf-training-example.py b/hf-training-example.py
index deb8cd3..cd518bb 100644
--- a/hf-training-example.py
+++ b/hf-training-example.py
@@ -4,6 +4,9 @@ import pandas as pd
 from torch.utils.data import Dataset, random_split
 from transformers import TrainingArguments, Trainer
 
+# # to save memory use bfloat16 on cpu
+# torch.set_default_dtype(torch.bfloat16)
+
 MODEL = 'decapoda-research/llama-7b-hf'
 DATA_FILE_PATH = 'datasets/elon_musk_tweets.csv'
 OUTPUT_DIR = './trained'
@@ -45,7 +48,7 @@ training_args = TrainingArguments(
     logging_dir='./logs',
     output_dir=OUTPUT_DIR,
     no_cuda=True,
-    # bf16=True,
+    bf16=True,
     per_device_eval_batch_size=1,
     per_device_train_batch_size=1)