fix the bug when using random erasing

ffiirree · Nov 27, 2021 · 899e7b9 · 899e7b9
1 parent b61b11f
commit 899e7b9
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/cvm/utils/factory.py b/cvm/utils/factory.py
@@ -291,12 +291,13 @@ def create_transforms(
         elif augment == 'autoaugment':
             ops.append(T.AutoAugment(autoaugment_policy))
 
-        if random_erasing > 0.0:
-            ops.append(T.RandomErasing(random_erasing))
-
     ops.append(T.PILToTensor())
     ops.append(T.ConvertImageDtype(torch.float))
     ops.append(T.Normalize(mean, std))
+
+    if is_training and random_erasing > 0.0:
+        ops.append(T.RandomErasing(random_erasing))
+
     return T.Compose(ops)