device arg wasn't removed from PrefetcherCuda instantiation of RE

pull/1239/head
Ross Wightman 3 years ago
parent 80ca078aed
commit d9b0b3d60f

@ -34,7 +34,7 @@ class PrefetcherCuda:
self.std = None self.std = None
if re_prob > 0.: if re_prob > 0.:
self.random_erasing = RandomErasing( 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: else:
self.random_erasing = None self.random_erasing = None

Loading…
Cancel
Save