order

pytorch · Apr 25, 2024 · ab787b3 · ab787b3
1 parent 37f59c5
commit ab787b3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -219,8 +219,6 @@ def loss_fn(pred, labels):
 
     metric_logger = build_metric_logger(job_config)
 
-    train_state = TrainState()
-
     if job_config.training.compile:
         if (
             job_config.activation_checkpoint.mode == "selective"
@@ -232,6 +230,8 @@ def loss_fn(pred, labels):
             )
         logger.info(f"Compiling each TransformerBlock with torch.compile")
 
+    train_state = TrainState()
+
     # train loop
     model.train()