Update train.py

train.py

@@ -145,9 +145,9 @@ wandb.init(
         "valid_set_size": len(valid_dataset),
         "model_parameters": parameters_count,
         "vocab_size": vocab_size,
-        "d_model": 384,
-        "d_inner": 768,
-        "n_layers": 6,
+        "d_model": model.d_model,
+        "d_inner": model.d_inner,
+        "n_layers": model.n_layers,
         "device": str(device),
         "model_name": model_name  # Log model_name
     }