Fixed bug with dataset argument

epidemicsound · Sep 1, 2023 · 0f9e3c8 · 0f9e3c8
1 parent 3648810
commit 0f9e3c8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/train_head.py b/src/train_head.py
@@ -68,7 +68,7 @@ def main(args):
     ############################
     # dataset
     ############################
-    dataset = DATASETS[backbone_args.dataset]
+    dataset = DATASETS[args.dataset]
     train_dataset = dataset(subset="train", transforms=transforms)
     val_dataset = dataset(subset="valid", transforms=transforms)
     if "nsynth" in args.dataset: