Minor fix to previous commit

cybershiptrooper · Sep 15, 2024 · 0f3369f · 0f3369f
1 parent caf1ff6
commit 0f3369f
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/iit/model_pairs/probed_sequential_pair.py b/iit/model_pairs/probed_sequential_pair.py
@@ -83,7 +83,8 @@ def train(
         test_set: IITDataset,
         epochs: int = 1000,
         use_wandb: bool = False,
-        wandb_name_suffix: str = "",
+        wandb_project: str = "iit",
+        wandb_name: str = "",
         optimizer_kwargs: dict = {},
     ) -> None:
         training_args = self.training_args
@@ -111,7 +112,7 @@ def train(
         loss_fn = t.nn.CrossEntropyLoss()
 
         if use_wandb and not wandb.run:
-            wandb.init(project="iit")
+            wandb.init(project=wandb_project, name=wandb_name)
 
         if use_wandb:
             wandb.config.update(training_args)