From d9b0b3d60f5a49403799886583b58ff2bd9d8b43 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Thu, 11 Nov 2021 16:21:13 -0800 Subject: [PATCH] device arg wasn't removed from PrefetcherCuda instantiation of RE --- timm/data/prefetcher_cuda.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timm/data/prefetcher_cuda.py b/timm/data/prefetcher_cuda.py index 0b36c027..3ae52dc2 100644 --- 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