rzzn
/
Mamba_test
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

更新 Mamba/mamba-main/train.py

main
rzzn 2024-07-30 17:53:00 +08:00
parent 9eea6c07af
commit f65b091fac
1 changed files with 99 additions and 99 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
Mamba/mamba-main/train.py
View File

@ -4,7 +4,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments,
from trl import SFTTrainer
from peft import LoraConfig
from datasets import Dataset
import torch
# 设置环境变量来避免内存碎片化
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"