fix augmix bug
@@ -352,7 +352,7 @@ class AugMix(object):
np.random.beta(self.aug_prob_coeff, self.aug_prob_coeff))
# image = Image.fromarray(image)
- mix = np.zeros([image.shape[1], image.shape[0], 3])
+ mix = np.zeros(image.shape)
for i in range(self.mixture_width):
image_aug = image.copy()
image_aug = Image.fromarray(image_aug)
