Another one

examples/pytorch/image-pretraining/run_mae.py

@@ -163,6 +163,13 @@ class CustomTrainingArguments(TrainingArguments):
         default=1e-3, metadata={"help": "Base learning rate: absolute_lr = base_lr * total_batch_size / 256."}
     )
 
+    def __post_init__(self):
+        # Compute absolute learning rate while args are mutable
+        if self.base_learning_rate is not None:
+            total_train_batch_size = self.train_batch_size * self.gradient_accumulation_steps * self.world_size
+            self.learning_rate = self.base_learning_rate * total_train_batch_size / 256
+        super().__post_init__()
+
 
 def collate_fn(examples):
     pixel_values = torch.stack([example["pixel_values"] for example in examples])
@@ -353,13 +360,6 @@ def preprocess_images(examples):
         # Set the validation transforms
         ds["validation"].set_transform(preprocess_images)
 
-    # Compute absolute learning rate
-    total_train_batch_size = (
-        training_args.train_batch_size * training_args.gradient_accumulation_steps * training_args.world_size
-    )
-    if training_args.base_learning_rate is not None:
-        training_args.learning_rate = training_args.base_learning_rate * total_train_batch_size / 256
-
     # Initialize our trainer
     trainer = Trainer(
         model=model,