Update vit.py

basedformer/models/vit.py

@@ -138,7 +138,7 @@ class VisionTransformer(base_image.BaseVisionModel):
             'activation': gelu_new,
             'image_size': (224, 224),
             'eps': 1e-5,
-            'device': torch.device('gpu'),
+            'device': torch.device('cuda'),
             'dtype': torch.float16,
         }
         super().__init__(self.default_config)