device arg wasn't removed from PrefetcherCuda instantiation of RE

3 years ago · d9b0b3d60f
parent 80ca078aed
commit d9b0b3d60f
1 changed files with 1 additions and 1 deletions
--- a/timm/data/prefetcher_cuda.py
+++ b/timm/data/prefetcher_cuda.py
@ -34,7 +34,7 @@ class PrefetcherCuda:
            self.std = None
        if re_prob > 0.:
            self.random_erasing = RandomErasing(
-                probability=re_prob, mode=re_mode, count=re_count, num_splits=num_aug_splits, device=device)
+                probability=re_prob, mode=re_mode, count=re_count, num_splits=num_aug_splits)
        else:
            self.random_erasing = None