vit model does not get imported automatically

megatron/model/__init__.py

@@ -30,7 +30,6 @@ def import_layernorm(fp32_residual_connection):
 
 
 from .distributed import *
-from .vit_model import VitModel
 from .bert_model import (BertModel,
                          BertModelFirstStage,
                          BertModelIntermediateStage,

pretrain_vit.py

@@ -19,7 +19,7 @@ import torch
 import torch.nn.functional as F
 from megatron import get_args, get_timers, mpu, print_rank_0
 from megatron.data.vit_dataset import build_train_valid_datasets
-from megatron.model import VitModel
+from megatron.model.vit_model import VitModel
 from megatron.training import pretrain
 from megatron.utils import average_losses_across_data_parallel_group
 

tasks/vision/classification.py

@@ -17,7 +17,7 @@
 
 from megatron import get_args
 from megatron import print_rank_0
-from megatron.model import VitModel
+from megatron.model.vit_model import VitModel
 from megatron.data.vit_dataset import build_train_valid_datasets
 from tasks.vision.eval_utils import accuracy_func_provider
 from tasks.vision.finetune_utils import finetune