Fix bugs that Mixup does not work device=cpu

3 years ago · 047a5ec05f
parent cd34913278
commit 047a5ec05f
1 changed files with 1 additions and 1 deletions
--- a/timm/data/mixup.py
+++ b/timm/data/mixup.py
@ -214,7 +214,7 @@ class Mixup:
            lam = self._mix_pair(x)
        else:
            lam = self._mix_batch(x)
-        target = mixup_target(target, self.num_classes, lam, self.label_smoothing)
+        target = mixup_target(target, self.num_classes, lam, self.label_smoothing, x.device)
        return x, target