diff --git a/train_diffusion_prior.py b/train_diffusion_prior.py
index c513c12..016b4e1 100644
--- a/train_diffusion_prior.py
+++ b/train_diffusion_prior.py
@@ -342,7 +342,7 @@ def main():
           RESUME,
           DPRIOR_PATH,
           config,
-          atgs.wandb_entity,
+          args.wandb_entity,
           args.wandb_project,
           args.learning_rate,
           args.max_grad_norm,