fix CutMix and RandomCutMix alpha arguments in get_mixup_cutmix in re…

…ferences/classification/transforms.py
pytorch · May 28, 2024 · 561921c · 561921c
1 parent c585a51
commit 561921c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/references/classification/transforms.py b/references/classification/transforms.py
@@ -19,9 +19,9 @@ def get_mixup_cutmix(*, mixup_alpha, cutmix_alpha, num_classes, use_v2):
         )
     if cutmix_alpha > 0:
         mixup_cutmix.append(
-            transforms_module.CutMix(alpha=mixup_alpha, num_classes=num_classes)
+            transforms_module.CutMix(alpha=cutmix_alpha, num_classes=num_classes)
             if use_v2
-            else RandomCutMix(num_classes=num_classes, p=1.0, alpha=mixup_alpha)
+            else RandomCutMix(num_classes=num_classes, p=1.0, alpha=cutmix_alpha)
         )
     if not mixup_cutmix:
         return None