Skip to content

Commit ebaec3c

Browse files
committed
fix axolotl training args dataclass annotation
1 parent 73e70e3 commit ebaec3c

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

Diff for: src/axolotl/utils/trainer.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import math
66
import os
77
import sys
8-
from dataclasses import field
8+
from dataclasses import dataclass, field
99
from pathlib import Path
1010
from typing import Optional
1111

@@ -29,6 +29,7 @@
2929
LOG = logging.getLogger("axolotl")
3030

3131

32+
@dataclass
3233
class AxolotlTrainingArguments(TrainingArguments):
3334
"""
3435
Extend the base TrainingArguments for axolotl helpers
@@ -188,7 +189,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
188189
if cfg.save_safetensors:
189190
training_arguments_kwargs["save_safetensors"] = cfg.save_safetensors
190191

191-
training_args = AxolotlTrainingArguments(
192+
training_args = AxolotlTrainingArguments( # pylint: disable=unexpected-keyword-arg
192193
per_device_train_batch_size=cfg.micro_batch_size,
193194
per_device_eval_batch_size=cfg.eval_batch_size
194195
if cfg.eval_batch_size is not None

0 commit comments

Comments
 (0)