diff --git a/bytelatent/train.py b/bytelatent/train.py
index f8761c7..0ee87df 100644
--- a/bytelatent/train.py
+++ b/bytelatent/train.py
@@ -130,7 +130,9 @@ def validate_train_args(args: TrainArgs, output_size: int):
     if args.model is not None:
         logger.info(f"Setting model output size to {args.model.vocab_size}")
         args.model.vocab_size = output_size
-        assert(args.model.max_encoder_seq_length == args.data.max_encoder_seq_length)
+        assert (
+            args.model.max_encoder_seq_length == args.data.max_encoder_seq_length
+        ), "max_encoder_seq_length for model and data should match"
 
     if args.entropy_model is not None:
         logger.info(f"Setting model output size to {args.entropy_model.vocab_size}")