Commit e2a4d6b1 authored by Arda Cihaner's avatar Arda Cihaner

Update vit.py

parent a05caa11
...@@ -138,7 +138,7 @@ class VisionTransformer(base_image.BaseVisionModel): ...@@ -138,7 +138,7 @@ class VisionTransformer(base_image.BaseVisionModel):
'activation': gelu_new, 'activation': gelu_new,
'image_size': (224, 224), 'image_size': (224, 224),
'eps': 1e-5, 'eps': 1e-5,
'device': torch.device('gpu'), 'device': torch.device('cuda'),
'dtype': torch.float16, 'dtype': torch.float16,
} }
super().__init__(self.default_config) super().__init__(self.default_config)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment