gpu fix

basedformer/models/vit.py

@@ -138,8 +138,8 @@ class VisionTransformer(base_image.BaseVisionModel):
             'activation': gelu_new,
             'image_size': (224, 224),
             'eps': 1e-5,
-            'device': torch.device('cpu'),
-            'dtype': torch.float32,
+            'device': torch.device('gpu'),
+            'dtype': torch.float16,
         }
         super().__init__(self.default_config)
         self.embed = ViTEmbeds(self.config)