@@ -1,7 +1,7 @@
 import unittest
 from super_gradients.training.utils.utils import HpmStruct
 from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationTestDatasetInterface
-from super_gradients import SgModel
+from super_gradients import Trainer
 from super_gradients.training.metrics import Accuracy, Top5
 
 
@@ -21,10 +21,10 @@ class TestViT(unittest.TestCase):
         """
         Validate vit_base
         """
-        model = SgModel("test_vit_base", device='cpu')
-        model.connect_dataset_interface(self.dataset, data_loader_num_workers=8)
-        model.build_model('vit_base', load_checkpoint=False)
-        model.train(training_params=self.train_params)
+        trainer = Trainer("test_vit_base", device='cpu')
+        trainer.connect_dataset_interface(self.dataset, data_loader_num_workers=8)
+        trainer.build_model('vit_base', load_checkpoint=False)
+        trainer.train(training_params=self.train_params)
 
 
 if __name__ == '__main__':