minor fixes

megatron/data/data_samplers.py

@@ -164,7 +164,7 @@ class MegatronPretrainingRandomSampler:
         current_epoch_samples = self.consumed_samples % active_total_samples
         assert current_epoch_samples % self.micro_batch_times_data_parallel_size == 0
 
-        if isinstance(dataset, RandomSeedDataset):
+        if isinstance(self.dataset, RandomSeedDataset):
             self.dataset.set_epoch(self.epoch)
 
         # data sharding and random sampling

megatron/data/vit_dataset.py

@@ -35,7 +35,7 @@ class ClassificationTransform():
                 T.ColorJitter(0.4, 0.4, 0.4, 0.1),
                 ImageNetPolicy(),
                 T.ToTensor(),
-                T.Normalize(*self.mean_std),
+                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
                 T.ConvertImageDtype(self.data_type)
             ])
         else:

megatron/model/vision/classification.py

@@ -19,8 +19,6 @@ import torch
 from megatron import get_args
 from megatron.model.utils import get_linear_layer
 from megatron.model.vision.vit_backbone import VitBackbone, VitMlpHead
-from megatron.model.vision.mit_backbone import mit_b3_avg
-from megatron.model.vision.utils import trunc_normal_
 from megatron.model.module import MegatronModule
 
 class VitClassificationModel(MegatronModule):