Update WANDB configs

VainF · Jul 7, 2024 · 4a24666 · 4a24666
1 parent a11bef3
commit 4a24666
Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/convnext_train.py b/convnext_train.py
@@ -41,8 +41,8 @@
 from timm.utils import ApexScaler, NativeScaler
 
 import torch_pruning as tp
-import pruning_bench as pbench
-pbench.pruning.patch_timm_forward()
+import pbench
+pbench.forward_patch.patch_timm_forward()
 
 from PIL import ImageFile
 ImageFile.LOAD_TRUNCATED_IMAGES = True
@@ -512,6 +512,7 @@ def main():
         wandb.init(
             # set the wandb project where this run will be logged
             project="Pruning", 
+            name=args.model.replace('/', '_'),
             # track hyperparameters and run metadata
             config=args,
         )

diff --git a/train.py b/train.py
@@ -263,7 +263,7 @@ def main(args):
     if torch.distributed.get_rank() == 0:
         wandb.init(
             # set the wandb project where this run will be logged
-            project="pbench", 
+            project="Pruning", 
             name=args.model.replace('/', '_'),
             # track hyperparameters and run metadata
             config=args,